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

Spot the error·Medium·4.0 · 0·~2 min

Attempt it

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

Mark at least one word to submit.

**The engineer's argument is wrong** because the causal mask and the padding mask block different things, and neither implies the other. The **causal mask** says: 'every query can attend to positions 0 through t-1, never to the future'. It is lower-triangular. It permits the past. The **padding mask** says: 'no query can attend to positions that are pad tokens'. It blocks specific columns regardless of their time index. In a right-padded batch, the pad tokens sit at the end of each short sequence. Consider a sequence with 3 real tokens followed by 5 pad tokens, batched alongside a sequence with 8 real tokens. The causal mask alone permits the 3rd real token to attend to all earlier positions, but in this short sequence those earlier positions are entirely real, so this case looks fine. The bug shows up when you consider what happens at the pad positions themselves. Position 4 (the first pad) is allowed to attend to positions 0-3 under the causal mask. That is fine because position 4 is going to be loss-masked anyway and its output does not matter. But position 5 (the second pad) is also allowed to attend backward... and position 7 is the same. The model still computes attention at every pad position, contaminates the K, V cache, and any code path that uses representations at pad positions reads contaminated data. **Loss masking** does not fix this because loss masking handles outputs, not inputs. The contamination has already happened in the attention layer. The fix is to apply the padding mask alongside the causal mask: combine them via logical AND (or sum the -inf masks) before softmax.

**The causal mask and the padding mask are orthogonal constraints,** not redundant. Each blocks something the other does not, so removing either one leaves a hole. - **Causal mask** blocks attention to **future** positions. Lower-triangular pattern over (time x time). - **Padding mask** blocks attention to **pad** positions. Column-based pattern over (any-time x pad-columns). In a right-padded batch with a short sequence, the causal mask permits attention from position t to positions 0..t-1. Some of those earlier positions may themselves be pad if the sequence is shorter than t suggests, but in right-padding the pads come after the reals, so this particular ordering of pads relative to reals does not immediately bite. **Where it does bite:** every position, including pad positions, still has its attention computed during training. The pad positions' K and V are mixed into the weighted sum at every later query position that attends to them under the causal mask. Even worse, in many implementations, the model is asked to predict next tokens at every position, including pad positions; the loss is then masked to ignore those, but the attention contamination has already corrupted the activations that produced those predictions. **The fix.** Combine the masks before softmax: ```python additive = causal_mask + padding_mask # both -inf where blocked scores = scores + additive weights = softmax(scores, dim=-1) ``` Any position that is either future or padding gets -inf and contributes nothing to the softmax. Loss masking is still needed to handle the prediction layer correctly, but it is a separate fix from the attention-level masking. **Why the engineer's intuition is tempting but wrong.** 'Pad tokens are in the past' is true, but the causal mask permits the past; it does not exclude any subset of the past. The causal mask is one-sided in time; the padding mask is one-sided in token-identity. Neither subsumes the other.

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

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

Detailed answer & concept explanation

~7 min read

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example.

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?

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

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

Don't say this

Red flags and common mistakes that signal junior thinking. Click to expand.

Most common mistake

Believing the causal mask subsumes the padding mask because 'pad tokens are in the past'. The causal mask permits the past; it does not exclude padding from the past.

60-second cheat sheet

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

References

Primary sources. Skim if you want the original framing.