A batch packs a causal LM with right-padded sequences and applies only the causal mask. Spot the mistake.

Attempt it

Click any words you think contain an error. Click again to unmark.

Mark at least one word to submit.

Quick reference

TL;DR

Causal mask blocks futures, padding mask blocks pad. Neither is a superset of the other. With right-padding you need both, applied together before softmax. Loss masking on outputs does not fix attention contamination.

Memory aid

Sign in for the mnemonic that makes this stick. See Pro pricing.

Easy to grasp

Picture a study group where some students arrived late and have empty notebook pages for the early part of the lecture. The class rule is 'only look at notes from earlier in the lecture' (causal mask). That rule says nothing about ignoring the empty pages, so a student trying to review their own notes from later in the lecture will happily read the empty pages from the late-arriving student next to them and try to study from blank paper. You need two rules together: 'only look at earlier notes' AND 'skip empty pages'. Right-padding plus causal mask only is one rule when you need both, and the model trains on blank-paper context.

Key concepts

Causal mask Padding mask Mask combination Loss masking

Quick scan · 2–4 min

Everything important, quickly.

Walk through what each mask blocks, show that the two constraints are orthogonal, demonstrate the bug with a worked example of a short right-padded sequence, explain why loss masking is a separate concern, and close with the FlashAttention-2 cu_seqlens replacement and a unit-test recipe.

Code example

python

import torch
import torch.nn.functional as F

B, T, n_heads, d_head = 2, 8, 4, 16
real_lengths = torch.tensor([3, 8])  # Seq 0: 3 real + 5 pad. Seq 1: all real.

Q = torch.randn(B, n_heads, T, d_head)
K = torch.randn(B, n_heads, T, d_head)
V = torch.randn(B, n_heads, T, d_head)

# Padding mask: -inf at pad columns, 0 at real columns. Shape (B, 1, 1, T).
is_real = torch.arange(T)[None, :] < real_lengths[:, None]  # (B, T)
pad_mask = torch.where(is_real, 0.0, float('-inf'))[:, None, None, :]

# Causal mask: -inf above the diagonal, 0 on and below. Shape (1, 1, T, T).
causal_mask = torch.triu(torch.full((T, T), float('-inf')), diagonal=1)
causal_mask = causal_mask[None, None, :, :]

# Combine BOTH masks before softmax. This is the correct approach.
scores = Q @ K.transpose(-2, -1) / (d_head ** 0.5)
scores = scores + causal_mask + pad_mask  # any position future OR pad becomes -inf
weights = F.softmax(scores, dim=-1)
output = weights @ V

# Loss masking is a SEPARATE step on the LM head output, not a replacement
# for the attention-level padding mask.

Real-world examples

Real products, models, and research that use this idea.

PyTorch's torch.nn.MultiheadAttention takes both attn_mask (causal) and key_padding_mask separately because they compose.
Hugging Face Transformers' models combine the attention_mask (padding) with an internally constructed causal mask in decoder layers.
FlashAttention-2's cu_seqlens API replaces the explicit (causal + padding) composition with per-example tiling, used in Llama 4 Maverick and Qwen 3.5 training.
vLLM's continuous batching avoids the issue entirely at inference by never having pad tokens in the batch.
Megatron-LM applies both causal and padding masks via a fused -inf addition before softmax during pretraining.

Interviewer follow-ups

What an interviewer would ask next. Try answering before peeking at the approach.

QWould left-padding instead of right-padding fix the bug?

Sign in or go Pro to see the approach for this follow-up.

QWhy does FlashAttention-2's cu_seqlens API make this entire discussion obsolete?

Sign in or go Pro to see the approach for this follow-up.

QSuppose you only have the padding mask and forget the causal mask in a decoder LM training loop. What happens?

Sign in or go Pro to see the approach for this follow-up.

60-second cheat sheet

The night-before-the-interview bullets. Scan these on the way to the call.

Sign in for 5 unlocks/day or go Pro for unlimited revision sheets. See Pro pricing.

References

Primary sources. Skim if you want the original framing.