The mask in attention is added to the scaled score matrix QK^T / sqrt(d_k) just before softmax. Masked positions get -inf in the score; softmax exponentiates that to 0, so those positions contribute exactly zero weight to the final attention distribution. Unmasked positions stay unchanged because the mask adds 0 to their scores.
The other options are wrong for specific reasons:
- **On the input embeddings**: too early. The embedding goes through Q, K, V projections regardless of whether the position is masked. Zeroing the input doesn't prevent attention from seeing the masked positions later.
- **On the post-softmax weights**: too late and breaks math. After softmax, the row sums to 1. If you zero out some weights, the row no longer sums to 1, and you'd need to renormalize the remaining weights. Doing the masking before softmax handles this automatically (softmax renormalizes the unmasked positions).
- **On the output after W_O**: way too late. By then the bad attention weights have already been used to compute the context vectors; zeroing the output positions doesn't undo that.
The additive before softmax pattern is standard across every transformer codebase. It's the cleanest mathematical placement and the one that fused kernels like FlashAttention expect.
**The mask is added to the scaled score matrix just before softmax.** Masked positions get -inf; softmax sends -inf to exactly 0, and unmasked positions renormalize over the remaining (nonzero) exponentials.
The formula in pseudocode:
```
scores = (Q @ K.transpose) / sqrt(d_k)
scores = scores + mask # mask has 0 at allowed, -inf at forbidden
weights = softmax(scores, dim=-1)
output = weights @ V
```
**Why this placement and not the alternatives:**
- **Before Q/K/V projection (on embeddings)**: too early. The projections still run regardless; attention still computes scores for masked positions; the mask has no effect.
- **After softmax (on weights)**: too late. Softmax already produced a valid probability distribution. If you zero out some weights post-softmax, the row no longer sums to 1, and the weighted-sum output is biased low for the unmasked positions. You'd need to renormalize manually, which is extra work and numerically uglier.
- **After W_O output**: far too late. The bad weights have already been baked into the context vectors, the output projection W_O has already mixed them across the head dimension, and zeroing now produces nonsense.
**The additive-with-(-inf) pattern is mathematically clean.** It uses softmax's own normalization to produce a valid distribution over the unmasked positions. The unmasked weights still sum to 1, the masked weights are exactly 0, and no extra renormalization is needed.
Dtype caveat: in fp16 or bf16, literal -inf can propagate to NaN through some accumulator paths. The portable idiom is `torch.finfo(dtype).min`, which is roughly -65504 in fp16 and -3.4e38 in fp32, both small enough to exponentiate to underflow.
Quick reference
TL;DR
The mask is added to the scaled score matrix just before softmax; masked positions get -inf, which softmax converts to exactly 0 weight.
Imagine a vote where every candidate's score gets converted into a vote percentage by a rule that exponentiates first. To truly disqualify a candidate, you have to push their score down to 'negative infinity bad', which is the only score the rule converts to zero votes. A score of zero, surprisingly, still gives them about one vote out of N. So the disqualification has to happen at the right step (before counting) using the right value (negative infinity), or the disqualified candidate still ends up with influence.
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.
Architectures that don't use softmax attention (linear attention, kernel attention, state-space models) These have different normalization mechanisms or none at all; the additive-with-(-inf) pattern is specific to softmax-based attention. Other architectures need different masking strategies.
Custom attention variants where the score matrix is never fully materialized FlashAttention computes attention block by block; you cannot directly add a full mask matrix. You either pass a function that produces the mask on demand or rely on the kernel's built-in causal flag.
Bidirectional encoder models (BERT, ViT) for tasks that don't need causal masking These use only padding masks (or no mask at all); the placement is the same but the mask content differs. Applying a causal mask here would be wrong.
Mixed-precision training with very small mask values that risk underflow to zero in fp16 Using -1e9 in fp16 overflows; use torch.finfo(dtype).min instead, which is dtype-safe. The placement is unchanged but the mask value matters.
Quick scan · 2–4 min
Everything important, quickly.
4m: why the mask must precede softmax, additive vs multiplicative masking, why -inf is the correct value, dtype-safe mask values, how causal and padding masks differ, and how FlashAttention fuses masking with softmax.
Code example
python
import torch
import math
B, H, T, D = 1, 1, 4, 8
q = torch.randn(B, H, T, D)
k = torch.randn(B, H, T, D)
v = torch.randn(B, H, T, D)
# Build causal mask: 1 at allowed (i >= j), 0 at forbidden
causal_mask = torch.tril(torch.ones(T, T))
# Compute scaled scores
scores = (q @ k.transpose(-2, -1)) / math.sqrt(D)
# Apply mask BEFORE softmax with -inf at forbidden positions
scores = scores.masked_fill(causal_mask == 0, float('-inf'))
# Softmax produces a valid distribution over unmasked positions
weights = torch.softmax(scores, dim=-1)
# Verify: row 0 attends only to col 0
assert weights[..., 0, 1:].sum().item() == 0.0
# Output
out = weights @ v
Real-world examples
Real products, models, and research that use this idea.
1PyTorch torch.nn.functional.scaled_dot_product_attention with is_causal=True applies the additive -inf mask internally before softmax.
2FlashAttention 2 and 3 fuse the mask application with the softmax computation, never materializing the full mask matrix.
3HuggingFace transformers uses torch.finfo(dtype).min as the mask value for dtype safety across fp16/bf16/fp32.
4Karpathy's nanoGPT uses the canonical pattern: scores.masked_fill(causal_mask == 0, float('-inf')) followed by softmax.
5vLLM and SGLang's attention kernels handle padding masks per-request and causal masks synthesized in-kernel for efficiency.
Interviewer follow-ups
What an interviewer would ask next. Try answering before peeking at the approach.
QHow would you detect a misplaced mask in production?
Sign in or go Pro to see the approach for this follow-up.
QWhy does FlashAttention support is_causal=True as a kernel-level flag instead of always passing a mask tensor?
Sign in or go Pro to see the approach for this follow-up.
QWhat's the right way to combine a causal mask with a padding mask?
Sign in or go Pro to see the approach for this follow-up.
Don't say this
Red flags and common mistakes that signal junior thinking. Click to expand.
Most common mistake
Applying the mask after softmax by zeroing weights. This breaks the row sum to 1 invariant and requires extra renormalization that fused kernels don't do for you.