PyTorch Attention Shapes: Q, K, V, Multi-Head Attention Masks and Transformer Dimension Errors
PyTorch: Zero to Hero — Step 07
Attention code is where tensor-shape mistakes stop being annoying and start becoming architectural.
A CNN usually makes its dimensional assumptions fairly obvious. Attention does not.
A tensor that starts as:
(batch, sequence, embedding)
is projected into Q, K and V, split into heads, transposed, multiplied, masked, normalized, multiplied again, transposed again, concatenated and projected back to the embedding dimension.
A single bad view, transpose, mask shape or head calculation can produce anything from an immediate runtime error to a model that trains while attending to the wrong tokens.
This post is a programmer-focused guide to making those transformations explicit.
We will build attention manually, inspect every shape, reproduce the common failures, and then compare the result with PyTorch’s higher-level attention APIs.
The shape contract to memorize
For most transformer code using batch_first=True, start with:
x: (B, T, C)
where:
B = batch size
T = sequence length
C = embedding dimension
For multi-head attention:
H = number of heads
D = head dimension
and therefore:
C = H * D
That identity is not optional.
If:
embed_dim = 512
num_heads = 8
then:
head_dim = embed_dim // num_heads
which gives:
64
The moment you split heads, your shape becomes:
(B, H, T, D)
That is the internal shape we will use throughout this article.
We can visualize the entire flow in one diagram:
flowchart TD
X["x: (B,T,C)"] --> LIN["Linear projections"]
LIN --> Q["Q: (B,T,C)"] & K["K: (B,T,C)"] & V["V: (B,T,C)"]
Q --> SH["split heads"]
K --> SH
V --> SH
SH --> QH["Q: (B,H,T,D)"] & KH["K: (B,H,T,D)"] & VH["V: (B,H,T,D)"]
QH --> SCORES["scores = Q @ K^T"]
KH --> SCORES
SCORES --> S["scores: (B,H,T,T)"]
S --> MASK["mask + softmax"]
MASK --> W["weights: (B,H,T,T)"]
W --> OUT["output = W @ V"]
VH --> OUT
OUT --> O1["out: (B,H,T,D)"]
O1 --> MERGE["merge heads"]
MERGE --> M1["(B,T,C)"]
M1 --> PROJ["output projection"]
PROJ --> Y["y: (B,T,C)"]
Start with a small executable example
import torch
B = 2
T = 5
C = 12
H = 3
D = C // H
x = torch.randn(B, T, C)
print('x:', x.shape)
print('head_dim:', D)
Output:
x: torch.Size([2, 5, 12])
head_dim: 4
Before doing anything else, assert the contract:
assert C % H == 0
Do this in production code too.
A surprising amount of transformer debugging starts with code that should have rejected an invalid configuration earlier.
Q, K and V are just projections
Create three linear projections:
import torch.nn as nn
q_proj = nn.Linear(C, C, bias=False)
k_proj = nn.Linear(C, C, bias=False)
v_proj = nn.Linear(C, C, bias=False)
q = q_proj(x)
k = k_proj(x)
v = v_proj(x)
print(q.shape)
print(k.shape)
print(v.shape)
All three are still:
(B, T, C)
For our example:
torch.Size([2, 5, 12])
Nothing attention-specific has happened yet.
The projections simply create three learned representations of the same input.
Split the embedding dimension into heads
We want:
(B, T, C)
to become:
(B, H, T, D)
The first step is reshape:
def split_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor:
B, T, C = x.shape
assert C % num_heads == 0
D = C // num_heads
x = x.reshape(B, T, num_heads, D)
x = x.transpose(1, 2)
return x
Now:
qh = split_heads(q, H)
kh = split_heads(k, H)
vh = split_heads(v, H)
print('Q:', qh.shape)
print('K:', kh.shape)
print('V:', vh.shape)
Output:
Q: torch.Size([2, 3, 5, 4])
K: torch.Size([2, 3, 5, 4])
V: torch.Size([2, 3, 5, 4])
That is:
(B, H, T, D)
Why transpose is necessary
After the reshape alone:
q.reshape(B, T, H, D)
we have:
(B, T, H, D)
But we want each head to perform attention independently across sequence positions.
So we move the head dimension before the sequence dimension:
q = q.reshape(B, T, H, D).transpose(1, 2)
Now:
(B, H, T, D)
That arrangement makes the matrix multiplication convenient.
The first classic bug: reshaping directly into the wrong order
This is wrong:
q_bad = q.reshape(B, H, T, D)
It has the desired shape.
But it does not have the desired meaning.
You did not move the head dimension.
You reinterpreted contiguous memory as if it were already arranged by head.
This is one of the nastiest tensor bugs because the shape check passes.
Correct:
q_good = q.reshape(B, T, H, D).transpose(1, 2)
The lesson is important:
Correct shape does not imply correct layout semantics.
Scaled dot-product attention from scratch
With:
Q: (B, H, T, D)
K: (B, H, T, D)
we want attention scores between every query position and every key position.
Transpose K’s final two dimensions:
scores = qh @ kh.transpose(-2, -1)
print(scores.shape)
Output:
torch.Size([2, 3, 5, 5])
The score matrix is:
(B, H, T, T)
Each head has a T x T matrix describing which token attends to which token.
Scale by the square root of the head dimension:
import math
scores = scores / math.sqrt(D)
Then softmax over the key dimension:
weights = torch.softmax(scores, dim=-1)
Finally multiply by V:
out = weights @ vh
print(out.shape)
Output:
torch.Size([2, 3, 5, 4])
Still:
(B, H, T, D)
A complete manual attention function
import math
import torch
def scaled_dot_product_attention(q, k, v, mask=None):
"""
q: (B, H, Tq, D)
k: (B, H, Tk, D)
v: (B, H, Tk, Dv)
"""
D = q.shape[-1]
scores = q @ k.transpose(-2, -1)
scores = scores / math.sqrt(D)
if mask is not None:
scores = scores.masked_fill(mask, float('-inf'))
weights = torch.softmax(scores, dim=-1)
output = weights @ v
return output, weights
For self-attention:
out, weights = scaled_dot_product_attention(qh, kh, vh)
print('out:', out.shape)
print('weights:', weights.shape)
Output:
out: torch.Size([2, 3, 5, 4])
weights: torch.Size([2, 3, 5, 5])
Cross-attention changes one dimension
Self-attention uses the same sequence for Q, K and V.
Cross-attention does not.
For cross-attention:
Q: (B, H, Tq, D)
K: (B, H, Tk, D)
V: (B, H, Tk, Dv)
The score tensor becomes:
(B, H, Tq, Tk)
That difference matters when constructing masks.
Do not hardcode square masks unless the operation is guaranteed to be self-attention.
Merge heads back together
After attention:
(B, H, T, D)
we need to return to:
(B, T, C)
Use:
def merge_heads(x: torch.Tensor) -> torch.Tensor:
B, H, T, D = x.shape
x = x.transpose(1, 2)
x = x.reshape(B, T, H * D)
return x
Test it:
merged = merge_heads(out)
print(merged.shape)
Output:
torch.Size([2, 5, 12])
Back to the original embedding dimension.
Why reshape is safer than view here
After:
x.transpose(1, 2)
the tensor may be non-contiguous.
This can break:
x.view(B, T, H * D)
A safer option is:
x.reshape(B, T, H * D)
Or explicitly:
x = x.transpose(1, 2).contiguous()
x = x.view(B, T, H * D)
When debugging transformer code, inspect:
print(x.is_contiguous())
print(x.stride())
The problem may not be the dimensions themselves.
It may be the memory layout after a transpose.
Build multi-head self-attention from scratch
import torch
import torch.nn as nn
import torch.nn.functional as F
class MultiHeadSelfAttention(nn.Module):
def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0):
super().__init__()
if embed_dim % num_heads != 0:
raise ValueError(
f'embed_dim={embed_dim} must be divisible by num_heads={num_heads}'
)
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False)
self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False)
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
self.dropout = dropout
def split_heads(self, x):
B, T, C = x.shape
x = x.reshape(B, T, self.num_heads, self.head_dim)
x = x.transpose(1, 2)
return x
def merge_heads(self, x):
B, H, T, D = x.shape
x = x.transpose(1, 2)
x = x.reshape(B, T, H * D)
return x
def forward(self, x, attn_mask=None, is_causal=False):
q = self.split_heads(self.q_proj(x))
k = self.split_heads(self.k_proj(x))
v = self.split_heads(self.v_proj(x))
y = F.scaled_dot_product_attention(
q,
k,
v,
attn_mask=attn_mask,
dropout_p=self.dropout if self.training else 0.0,
is_causal=is_causal,
)
y = self.merge_heads(y)
y = self.out_proj(y)
return y
Test it:
model = MultiHeadSelfAttention(embed_dim=12, num_heads=3)
x = torch.randn(2, 5, 12)
y = model(x)
print(y.shape)
Expected:
torch.Size([2, 5, 12])
The head divisibility error
This configuration is invalid:
MultiHeadSelfAttention(embed_dim=10, num_heads=3)
because:
10 / 3
is not an integer.
Do not rely on a later reshape to fail.
Reject it explicitly:
if embed_dim % num_heads != 0:
raise ValueError(...)
This produces a much better error than a mysterious reshape failure 40 lines later.
batch_first is a source of endless shape bugs
PyTorch’s nn.MultiheadAttention historically defaulted to sequence-first ordering:
(T, B, C)
When:
batch_first=True
it uses:
(B, T, C)
For most modern application code, I strongly prefer:
mha = nn.MultiheadAttention(
embed_dim=64,
num_heads=8,
batch_first=True,
)
Then:
x = torch.randn(4, 32, 64)
out, weights = mha(x, x, x)
print(out.shape)
print(weights.shape)
The output is:
(4, 32, 64)
If attention weights are averaged across heads, their shape is typically:
(4, 32, 32)
Debug batch_first explicitly
Write this helper:
def assert_btc(x, *, embed_dim=None):
assert x.ndim == 3, f'expected (B,T,C), got {tuple(x.shape)}'
B, T, C = x.shape
if embed_dim is not None:
assert C == embed_dim, (
f'expected embedding dimension {embed_dim}, got {C}'
)
return B, T, C
Then:
B, T, C = assert_btc(x, embed_dim=64)
Treat tensor shapes as interfaces.
Do not leave them as assumptions in your head.
Causal attention
A language model must not look at future tokens while predicting the current token.
For a sequence of length 5, conceptually the allowed attention pattern is:
1 0 0 0 0
1 1 0 0 0
1 1 1 0 0
1 1 1 1 0
1 1 1 1 1
A manual boolean mask that blocks the upper triangle can be built with:
T = 5
causal_mask = torch.triu(
torch.ones(T, T, dtype=torch.bool),
diagonal=1,
)
print(causal_mask)
For the manual implementation in this article, True means blocked because we use:
scores.masked_fill(mask, float('-inf'))
That convention is extremely common.
But do not assume every PyTorch attention API uses boolean masks with exactly the same meaning.
This is one of the most important debugging rules in attention code.
Mask semantics are API-specific
This is where experienced programmers still lose time.
For nn.MultiheadAttention, a boolean attn_mask uses True to mean the position is not allowed to attend.
For torch.nn.functional.scaled_dot_product_attention, a boolean attn_mask uses True to mean the position participates in attention.
Those are opposite boolean conventions.
So this:
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
may need inversion depending on which API you pass it to.
Do not write generic mask code without documenting the target API.
Prefer is_causal=True when appropriate
With scaled dot-product attention:
y = F.scaled_dot_product_attention(
q,
k,
v,
is_causal=True,
)
This avoids manually constructing a causal mask in many cases.
It also gives PyTorch more information about the operation.
For debugging, this is valuable because you eliminate one whole source of mask-shape and mask-semantics bugs.
Padding masks are different from causal masks
Suppose we batch variable-length sequences:
sequence 1: [A, B, C, D]
sequence 2: [E, F, PAD, PAD]
A padding mask marks tokens that should not be attended to.
For nn.MultiheadAttention:
key_padding_mask = torch.tensor([
[False, False, False, False],
[False, False, True, True ],
])
Shape:
(B, T)
Then:
out, weights = mha(
x,
x,
x,
key_padding_mask=key_padding_mask,
)
This is not the same object as an attention mask.
Attention mask vs padding mask
Think of them as answering different questions.
Attention mask
Can query position i attend to key position j?
Typical shape:
(Tq, Tk)
or sometimes a batch/head-expanded form.
Padding mask
Is this key position real data or padding?
Typical shape:
(B, Tk)
Mixing the two is a classic transformer bug.
Inspect masks like tensors, not metadata
Add a helper:
def describe_mask(name, mask):
if mask is None:
print(f'{name}: None')
return
print(
f'{name}: '
f'shape={tuple(mask.shape)} '
f'dtype={mask.dtype} '
f'device={mask.device} '
f'true={int(mask.bool().sum())}/{mask.numel()}'
)
Use it before attention calls:
describe_mask('attn_mask', attn_mask)
describe_mask('key_padding_mask', key_padding_mask)
This catches many bugs immediately.
The mask device bug
This fails when the model is on CUDA but the mask is on CPU:
x = x.cuda()
mask = torch.ones(T, T, dtype=torch.bool)
Then attention mixes devices.
Build masks on the same device:
mask = torch.ones(
T,
T,
dtype=torch.bool,
device=x.device,
)
Or:
mask = mask.to(x.device)
For reusable modules, derive the device from an input tensor rather than hardcoding 'cuda'.
The mask dtype bug
A mask may be boolean or floating-point depending on the API and intended behavior.
A floating mask usually contains additive bias values such as:
0
-inf
Example:
mask = torch.zeros(T, T)
mask = mask.masked_fill(
torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1),
float('-inf'),
)
Always print:
mask.dtype
when debugging attention.
Shape alone is not enough.
A reusable attention shape tracer
def trace_attention_shapes(q, k, v, mask=None):
print('Q', tuple(q.shape), q.dtype, q.device)
print('K', tuple(k.shape), k.dtype, k.device)
print('V', tuple(v.shape), v.dtype, v.device)
assert q.ndim >= 3
assert k.ndim >= 3
assert v.ndim >= 3
assert q.shape[-1] == k.shape[-1], (
f'Q/K head dimensions differ: {q.shape[-1]} vs {k.shape[-1]}'
)
assert k.shape[-2] == v.shape[-2], (
f'K/V sequence lengths differ: {k.shape[-2]} vs {v.shape[-2]}'
)
if mask is not None:
print('MASK', tuple(mask.shape), mask.dtype, mask.device)
Call it immediately before attention:
trace_attention_shapes(q, k, v, mask)
Check the score matrix directly
When attention is behaving strangely, stop hiding behind the module.
Compute:
scores = q @ k.transpose(-2, -1)
Then inspect:
print(scores.shape)
print(scores.min().item())
print(scores.max().item())
print(scores.mean().item())
print(torch.isfinite(scores).all().item())
After masking:
masked_scores = scores.masked_fill(mask, float('-inf'))
check:
print(torch.isfinite(masked_scores).any(dim=-1))
A row that is entirely -inf will produce invalid softmax behavior.
Entirely masked rows
This is a subtle and important failure mode.
If every key for a query position is masked:
[-inf, -inf, -inf, -inf]
then there is no valid distribution to normalize.
That can lead to NaNs or undefined behavior depending on the implementation and dtype.
Write a diagnostic:
def assert_not_fully_masked(mask):
if mask.dtype != torch.bool:
return
fully_masked = mask.all(dim=-1)
if fully_masked.any():
bad = fully_masked.nonzero(as_tuple=False)
raise ValueError(f'fully masked attention rows at {bad.tolist()}')
Adapt the logic to your API’s boolean convention.
Attention weights should sum to one
For a manual softmax implementation:
weights = torch.softmax(scores, dim=-1)
verify:
row_sums = weights.sum(dim=-1)
print(row_sums)
They should be approximately one for valid rows.
Assert it:
assert torch.allclose(
row_sums,
torch.ones_like(row_sums),
atol=1e-5,
)
This is a very useful unit test.
Verify causal attention directly
For causal attention, the upper triangle should have zero attention probability.
future = torch.triu(
torch.ones(T, T, dtype=torch.bool),
diagonal=1,
)
Given:
weights: (B, H, T, T)
check:
future_weights = weights[..., future]
print(future_weights.abs().max())
For a properly masked implementation, this should be effectively zero.
Compare manual attention with PyTorch SDPA
A very good debugging technique is to compare your implementation against scaled_dot_product_attention.
import torch.nn.functional as F
manual_out, _ = scaled_dot_product_attention(
qh,
kh,
vh,
)
pytorch_out = F.scaled_dot_product_attention(
qh,
kh,
vh,
dropout_p=0.0,
)
print(torch.allclose(manual_out, pytorch_out, atol=1e-5))
If this is false, inspect:
scale
mask convention
softmax dimension
Q/K transpose
head order
before blaming autograd or the optimizer.
nn.MultiheadAttention as a reference implementation
PyTorch also exposes:
mha = nn.MultiheadAttention(
embed_dim=C,
num_heads=H,
batch_first=True,
)
Use:
out, weights = mha(
x,
x,
x,
need_weights=True,
)
print(out.shape)
print(weights.shape)
If you need weights per head:
out, weights = mha(
x,
x,
x,
need_weights=True,
average_attn_weights=False,
)
print(weights.shape)
For batched input, this gives a shape like:
(B, H, T, T)
That is far more useful for debugging individual heads.
Disable weight return when you do not need it
For production code:
out, _ = mha(
x,
x,
x,
need_weights=False,
)
This can allow PyTorch to use a more optimized attention path.
Do not keep attention weights just because a tutorial did.
They can be expensive for long sequences.
Attention memory grows quadratically with sequence length
The score tensor shape is:
(B, H, T, T)
Its element count is:
B * H * T * T
Write a helper:
def attention_score_memory_mb(
batch_size,
num_heads,
seq_len,
bytes_per_element=4,
):
elements = batch_size * num_heads * seq_len * seq_len
return elements * bytes_per_element / (1024 ** 2)
Example:
for T in [128, 256, 512, 1024, 2048]:
mb = attention_score_memory_mb(
batch_size=8,
num_heads=8,
seq_len=T,
)
print(T, f'{mb:.1f} MB')
This is only the score tensor.
Training needs considerably more memory than this simple estimate.
A transformer block with explicit shapes
class TransformerBlock(nn.Module):
def __init__(
self,
embed_dim: int,
num_heads: int,
mlp_ratio: int = 4,
dropout: float = 0.0,
):
super().__init__()
self.norm1 = nn.LayerNorm(embed_dim)
self.attn = MultiHeadSelfAttention(
embed_dim=embed_dim,
num_heads=num_heads,
dropout=dropout,
)
self.norm2 = nn.LayerNorm(embed_dim)
hidden_dim = embed_dim * mlp_ratio
self.mlp = nn.Sequential(
nn.Linear(embed_dim, hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, embed_dim),
)
def forward(self, x, is_causal=True):
assert x.ndim == 3
x = x + self.attn(
self.norm1(x),
is_causal=is_causal,
)
x = x + self.mlp(self.norm2(x))
return x
Test:
block = TransformerBlock(
embed_dim=64,
num_heads=8,
)
x = torch.randn(4, 32, 64)
y = block(x)
print(y.shape)
Expected:
torch.Size([4, 32, 64])
A transformer block should preserve (B, T, C).
That is an excellent invariant to test.
Unit-test your transformer shapes
def test_transformer_block_preserves_shape():
block = TransformerBlock(
embed_dim=64,
num_heads=8,
)
x = torch.randn(3, 17, 64)
y = block(x)
assert y.shape == x.shape
Test several sequence lengths:
import pytest
@pytest.mark.parametrize('seq_len', [1, 2, 7, 32, 129])
def test_transformer_variable_sequence_length(seq_len):
block = TransformerBlock(64, 8)
x = torch.randn(2, seq_len, 64)
y = block(x)
assert y.shape == x.shape
Do not test only the sequence length from your training config.
Shape bugs often hide in boundary cases.
Test sequence length 1
Sequence length one is particularly useful:
x = torch.randn(2, 1, 64)
A surprising amount of attention code accidentally squeezes dimensions and breaks here.
Avoid unqualified:
x.squeeze()
Prefer explicit dimensions:
x.squeeze(dim=1)
when that is actually what you intend.
Test batch size 1
The same rule applies to:
x = torch.randn(1, 32, 64)
If your code behaves differently for batch size 1, look for:
squeeze()
indexing that removes dimensions
manual batch broadcasting
incorrect mask expansion
A practical attention debugging function
def diagnose_attention(
q,
k,
v,
*,
mask=None,
num_heads=None,
):
print('--- attention diagnostic ---')
for name, tensor in [('Q', q), ('K', k), ('V', v)]:
print(
f'{name}: '
f'shape={tuple(tensor.shape)} '
f'dtype={tensor.dtype} '
f'device={tensor.device} '
f'contiguous={tensor.is_contiguous()}'
)
if q.shape[-1] != k.shape[-1]:
print('ERROR: Q and K head dimensions differ')
if k.shape[-2] != v.shape[-2]:
print('ERROR: K and V sequence lengths differ')
if num_heads is not None:
if q.ndim >= 4 and q.shape[-3] != num_heads:
print(
f'WARNING: expected {num_heads} heads, '
f'Q has dimension {q.shape[-3]}'
)
if mask is not None:
print(
f'MASK: '
f'shape={tuple(mask.shape)} '
f'dtype={mask.dtype} '
f'device={mask.device}'
)
if mask.device != q.device:
print('ERROR: mask and Q are on different devices')
This is intentionally boring.
Boring diagnostics are good diagnostics.
Common error: mat1 and mat2 shapes cannot be multiplied
Suppose your attention output is:
(B, T, 256)
but your output projection expects:
nn.Linear(512, 512)
PyTorch eventually reports a matrix multiplication error.
The real bug happened earlier.
Print:
print('before out_proj:', y.shape)
print('weight:', self.out_proj.weight.shape)
The final dimension going into nn.Linear must match in_features.
Common error: wrong transpose on K
This is wrong:
scores = q @ k.transpose(1, 2)
when Q and K have:
(B, H, T, D)
You want to transpose:
T and D
which are the final two dimensions:
scores = q @ k.transpose(-2, -1)
Using negative dimension indices makes attention code much easier to reason about.
Common error: softmax over the wrong dimension
Correct:
weights = torch.softmax(scores, dim=-1)
The final dimension represents the key positions being selected for each query.
If you accidentally use:
softmax(scores, dim=-2)
you normalize across queries instead.
The code runs.
The model is wrong.
This is exactly why attention weights should be tested to ensure the expected axis sums to one.
Common error: forgetting the scale
This:
scores = q @ k.transpose(-2, -1)
should normally be scaled by:
scores = scores / math.sqrt(q.shape[-1])
Without scaling, score magnitudes grow with head dimension and softmax can become excessively sharp.
Again, the code runs.
The training dynamics change.
Common error: applying softmax before masking
Wrong:
weights = torch.softmax(scores, dim=-1)
weights = weights.masked_fill(mask, 0)
Now the remaining probabilities no longer necessarily sum to one.
Correct order:
scores = scores.masked_fill(mask, float('-inf'))
weights = torch.softmax(scores, dim=-1)
Mask logits first.
Normalize second.
Common error: mask broadcasts but across the wrong axes
Broadcasting can make incorrect code look valid.
Suppose scores are:
(B, H, T, T)
and your mask is accidentally:
(B, T)
Depending on the operation, broadcasting may fail or may align dimensions differently from what you intended.
When you want a batch padding mask to broadcast across heads and query positions, make the transformation explicit:
padding_mask = padding_mask[:, None, None, :]
Now:
(B, 1, 1, T)
which broadcasts clearly against:
(B, H, T, T)
Explicit broadcasting is easier to debug than clever broadcasting.
Combined causal and padding masks
For a manual implementation where True means blocked:
B, T = padding_mask.shape
causal = torch.triu(
torch.ones(T, T, dtype=torch.bool, device=x.device),
diagonal=1,
)
causal = causal[None, None, :, :]
padding = padding_mask[:, None, None, :]
combined = causal | padding
Shapes:
graph TD
CAUSAL["causal: (1,1,T,T)"] --> COMBINED[combined: (B,1,T,T)]
PAD["padding: (B,1,1,T)"] --> COMBINED
COMBINED --> BROADCAST["broadcasts over heads → (B,H,T,T)"]
causal: (1, 1, T, T)
padding: (B, 1, 1, T)
combined:(B, 1, T, T)
That broadcasts over heads.
Write the shapes down.
Do not trust yourself to remember them.
Validate mask broadcasting before training
scores = torch.randn(B, H, T, T)
try:
test = scores.masked_fill(combined, float('-inf'))
print('broadcasted shape:', test.shape)
except RuntimeError as exc:
print('mask broadcast failed:', exc)
Do this in a unit test rather than discovering the problem at step 80,000.
A tiny causal language-model attention test
B = 2
T = 8
C = 32
H = 4
attn = MultiHeadSelfAttention(C, H)
x = torch.randn(B, T, C)
y = attn(x, is_causal=True)
assert y.shape == x.shape
assert torch.isfinite(y).all()
This should be the minimum smoke test for every custom attention module.
Test gradients too
A forward pass is not enough.
attn = MultiHeadSelfAttention(32, 4)
x = torch.randn(2, 8, 32, requires_grad=True)
y = attn(x, is_causal=True)
loss = y.square().mean()
loss.backward()
assert x.grad is not None
assert torch.isfinite(x.grad).all()
for name, param in attn.named_parameters():
assert param.grad is not None, f'missing gradient: {name}'
assert torch.isfinite(param.grad).all(), f'bad gradient: {name}'
Attention code that only works in inference is not finished.
Compare gradients between manual and reference implementations
For low-level debugging, initialize two implementations with matching weights and compare:
forward outputs
gradients with respect to Q/K/V
parameter gradients
torch.autograd.grad is useful here:
grads = torch.autograd.grad(
outputs=loss,
inputs=(q, k, v),
)
If outputs match but gradients do not, you have found a deeper implementation problem.
Do not materialize attention weights unless you need them
This matters for memory.
If your application does not need attention maps:
out, _ = mha(
x,
x,
x,
need_weights=False,
)
Returning full attention matrices for long sequences can add meaningful memory and compute overhead.
Debug with weights when debugging.
Turn them off when you are done.
Shape contracts for a transformer stack
A useful transformer invariant is:
input: (B, T, C)
block: (B, T, C)
block: (B, T, C)
block: (B, T, C)
output: (B, T, C)
Only specialized boundaries should change that contract.
For example:
token IDs: (B, T)
embedding: (B, T, C)
transformer: (B, T, C)
LM head: (B, T, V)
where:
V = vocabulary size
This is the architecture we will assemble in Step 10.
Build a shape-audited transformer stack
class ShapeCheckedBlock(TransformerBlock):
def forward(self, x, is_causal=True):
input_shape = x.shape
y = super().forward(x, is_causal=is_causal)
if y.shape != input_shape:
raise RuntimeError(
f'transformer block changed shape: '
f'{tuple(input_shape)} -> {tuple(y.shape)}'
)
return y
You can remove these checks later if profiling proves they matter.
During development, they are cheap insurance.
Hooks for transformer debugging
You can reuse forward hooks from earlier posts:
def shape_hook(name):
def hook(module, inputs, output):
def shape_of(value):
if isinstance(value, torch.Tensor):
return tuple(value.shape)
return type(value).__name__
print(
name,
'in=', [shape_of(v) for v in inputs],
'out=', shape_of(output),
)
return hook
Register them selectively:
handles = []
for name, module in model.named_modules():
if isinstance(module, (nn.Linear, nn.LayerNorm, MultiHeadSelfAttention)):
handles.append(
module.register_forward_hook(shape_hook(name))
)
After debugging:
for handle in handles:
handle.remove()
Debug values as well as shapes
A tensor can have the perfect shape and still be broken.
Add:
def tensor_health(name, x):
print(
name,
'shape=', tuple(x.shape),
'mean=', x.mean().item(),
'std=', x.std().item(),
'min=', x.min().item(),
'max=', x.max().item(),
'finite=', torch.isfinite(x).all().item(),
)
Use it on:
Q
K
V
attention scores
attention weights
attention output
Attention failures are often numerical, not dimensional.
Detect suspiciously uniform attention
Uniform attention is not always wrong, especially at initialization.
But if it persists unexpectedly, inspect it.
entropy = -(weights * (weights + 1e-12).log()).sum(dim=-1)
print('attention entropy:', entropy.mean().item())
You can compare this over training.
Very high entropy means broadly distributed attention.
Very low entropy means highly concentrated attention.
Do not optimize this metric blindly.
Use it as an observation.
Detect collapsed heads
If multiple heads become nearly identical, compare them:
# weights: (B, H, T, T)
head_means = weights.mean(dim=(0, 2, 3))
print(head_means)
For deeper analysis, flatten each head’s attention map and compute similarity.
Again, this is diagnostic evidence, not a universal quality metric.
A compact debugging checklist
When attention code fails, check these in order:
flowchart TD
A["input shape (B,T,C)?"] -->|No| A1[Fix input pipeline]
A -->|Yes| B["C divisible by H?"]
B -->|No| B1[Adjust embed_dim or num_heads]
B -->|Yes| C["After split: (B,H,T,D)?"]
C -->|No| C1[Check reshape + transpose]
C -->|Yes| D["K transposed (-2,-1)?"]
D -->|No| D1[Fix transpose axes]
D -->|Yes| E["Softmax over dim=-1?"]
E -->|No| E1[Correct softmax axis]
E -->|Yes| F["Mask device, dtype, semantics correct?"]
F -->|No| F1[Align mask with API]
F -->|Yes| G["Any row fully masked?"]
G -->|Yes| G1[Check padding / causal logic]
G -->|No| H["Weights sum ≈1?"]
H -->|No| H1[Check mask before softmax]
H -->|Yes| I["After merge: (B,T,C)?"]
I -->|No| I1[Fix merge heads]
I -->|Yes| J["Output proj in_features matches?"]
J -->|No| J1[Verify C in out_proj]
J -->|Yes| K["Gradients finite?"]
K -->|No| K1[Check numerical stability]
K -->|Yes| L["All good"]
That list will solve a surprising percentage of transformer bugs.
The key mental model
Multi-head attention looks complicated because a lot of tensor transformations happen close together.
But the core shape flow is simple:
(B, T, C)
|
| linear projections
v
Q, K, V: (B, T, C)
|
| split heads
v
(B, H, T, D)
|
| Q @ K^T
v
scores: (B, H, Tq, Tk)
|
| mask + softmax
v
weights: (B, H, Tq, Tk)
|
| weights @ V
v
(B, H, Tq, D)
|
| merge heads
v
(B, Tq, C)
|
| output projection
v
(B, Tq, C)
If you can account for every dimension in that pipeline, attention stops being magic.
It becomes matrix multiplication with disciplined bookkeeping.
Challenge 1: break head splitting
Start with:
x = torch.randn(2, 8, 32)
Implement head splitting incorrectly with:
x.reshape(2, 4, 8, 8)
Then implement it correctly with reshape + transpose.
Compare values, not only shapes.
Explain why the first version is semantically wrong despite having the same dimensions.
Challenge 2: create a causal-mask bug
Build a boolean causal mask.
Use it once with your manual attention implementation and once with F.scaled_dot_product_attention.
Observe the difference in boolean-mask convention.
Then write a helper that makes the convention explicit in the function name.
For example:
def blocked_positions_to_sdpa_mask(blocked):
return ~blocked
The point is not the helper itself.
The point is to make semantics visible in code.
Challenge 3: test cross-attention
Create:
q = torch.randn(2, 4, 7, 16)
k = torch.randn(2, 4, 11, 16)
v = torch.randn(2, 4, 11, 16)
Predict the output shape before running anything.
Then verify it.
The answer should be:
(2, 4, 7, 16)
The score matrix should be:
(2, 4, 7, 11)
Challenge 4: instrument a transformer block
Add diagnostics that record:
input shape
Q shape
K shape
V shape
score shape
attention-weight shape
merged-head shape
output shape
Make the code raise as soon as one contract is violated.
Do not wait for a later matrix multiplication to fail.
Where the series goes next
We now have enough PyTorch machinery to build and debug real architectures.
The remaining stages are:
Step 08 — PyTorch Model Not Learning? A Systematic Debugging Guide
Step 09 — PyTorch Performance: CUDA Memory, Profiling and torch.compile
Step 10 — Build a Small Language Model From Scratch
Step 08 moves away from isolated APIs and into the question that eventually hits every machine-learning programmer:
The code runs. The loss exists. The optimizer steps. Why is the model still not learning?
That is where we go next.
