At which step in the attention pipeline does the mask actually take effect?

MCQ·Easy·4.0 · 0·~1 min

Attempt it

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.

Memory aid

Easy to grasp

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.

Key concepts

Softmax function Causal mask Padding mask

Detailed answer & concept explanation

~6 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.

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.

PyTorch torch.nn.functional.scaled_dot_product_attention with is_causal=True applies the additive -inf mask internally before softmax.
FlashAttention 2 and 3 fuse the mask application with the softmax computation, never materializing the full mask matrix.
HuggingFace transformers uses torch.finfo(dtype).min as the mask value for dtype safety across fp16/bf16/fp32.
Karpathy's nanoGPT uses the canonical pattern: scores.masked_fill(causal_mask == 0, float('-inf')) followed by softmax.
vLLM 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?

QWhy does FlashAttention support is_causal=True as a kernel-level flag instead of always passing a mask tensor?

QWhat's the right way to combine a causal mask with a padding mask?

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.

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.