Articles
Attention Mask Interview Questions: Causal Masks, Padding, and Cached Decoding
Practice attention mask interview questions with causal and padding ledgers, PyTorch boolean conventions, cached decoding offsets, and invariance tests.
A decoder works during full-sequence evaluation but changes its answers when a KV cache is enabled. The tensor shapes still match. One possible cause is a mask whose rows use local query indices while its columns represent the entire cached prefix.
Attention mask interview questions are answered more clearly by writing the allowed query-key pairs before choosing a tensor representation. Causal order, padding, boolean conventions, and cache positions are separate concerns. A correctly shaped mask can still express the wrong relation.
Evidence boundary: Current PyTorch 2.14 API documentation and Hugging Face cache documentation establish the software facts linked below. The mask ledgers and numerical tests are original exercises and engineering inference. No candidate reports or employer-specific interview frequencies are claimed. Start with Implement Masked Multi-Head Self-Attention and label each mask dimension.
What does the mask change in attention?
For query length L and key length S, each attention row compares one query with S keys. The mask controls which key scores can contribute to that row's normalization and weighted value sum.
Original mathematical convention: Write the score matrix as Q @ K.T / sqrt(d). Add zero to allowed scores and negative infinity to forbidden scores, then apply softmax along the key dimension. A boolean representation can encode the same relation, but its meaning depends on the API.
Multiplying forbidden logits by zero is not equivalent. Zero is still a finite score and receives positive probability after softmax. Likewise, masking the weighted output after normalization generally leaves the remaining weights improperly normalized.
An attention mask is not a loss mask. Ignoring padded targets in cross-entropy controls which training terms count; it does not stop real queries from reading padded keys. Conversely, blocking padded keys does not automatically remove padded query positions from the loss.
Engineering inference: Keep names such as allowed_pairs, valid_keys, and valid_targets distinct. Reusing a tensor named mask for all three operations makes an interface mismatch harder to see.
How do API boolean conventions differ?
Official API fact: In PyTorch scaled dot-product attention, a boolean attn_mask value of True means that a pair participates. A floating mask is added to scores. See scaled_dot_product_attention.
Official API fact: In MultiheadAttention, True in a boolean key_padding_mask means that the key is ignored; True in its boolean attn_mask marks a prohibited pair. See MultiheadAttention.
| Interface or representation | Meaning of True or value | Conversion from allowed pairs |
|---|---|---|
| SDPA boolean attention mask | Pair is allowed | Use allowed relation directly |
| MHA boolean attention mask | Pair is blocked | Invert allowed relation |
| MHA boolean key padding mask | Key is ignored | Invert key validity |
| Additive score bias | Zero allows; negative infinity blocks | Convert relation to score bias |
This distinction is semantic, not cosmetic. Inverting a correct mask can prevent every intended context token from contributing. Check the exact function, parameter, and installed version rather than assuming one framework-wide convention.
A key-validity vector also differs from a query-key matrix. The former selects keys for a batch item; the latter can express a different allowed set for each query. State which object you pass before discussing broadcasting.
What should the full-sequence causal mask look like?
Original exercise: Four real tokens occupy positions zero through three. Queries and keys represent those same positions. Let one mean allowed and zero mean forbidden:
key position
0 1 2 3
query 0 1 0 0 0
query 1 1 1 0 0
query 2 1 1 1 0
query 3 1 1 1 1
Each token may read itself and the preceding tokens. This diagonal convention matches ordinary decoder attention over shifted next-token targets. If a task forbids reading the current token, state that different objective and change the relation deliberately.
Original check: Alter a future token's key and value while holding the earlier queries and permitted prefix fixed. Earlier outputs should remain unchanged. This tests the causal relation, not merely the mask's triangular appearance.
For encoder bidirectional attention, future restriction may be inappropriate. For cross-attention, query and key positions can describe different sequences; applying a decoder's square triangle without explaining the intended alignment is not a general solution.
How do padding and causal restrictions combine?
Original exercise: A padded batch item has two real tokens followed by two padded keys. Under the same four-position causal layout, real query rows should allow:
query 0 1 0 0 0
query 1 1 1 0 0
The relation is the intersection of causal permission and valid-key permission. For an allowed-True representation, use logical AND. For a blocked-True representation, combine the corresponding blocked conditions with logical OR.
Masking keys does not define the fate of padded query outputs. You can ignore those outputs in the downstream loss or explicitly handle them according to the model contract. Do not accidentally turn them into meaningful predictions.
Left padding creates an empty-row case. An early padded query can have no valid allowed key after combining the conditions. A naive softmax of an all-negative-infinity row has an undefined normalization. Decide whether such rows are avoided, ignored, or given an explicit harmless handling policy.
Engineering inference: Do not replace a fully masked row with a uniform distribution and present it as valid attention. That invents access to forbidden values. Some optimized implementations have specific behavior for empty rows; verify your actual backend and ensure downstream code treats those outputs appropriately.
Also test batch items with different valid lengths. A single mask accidentally broadcast across the entire batch may look correct for the longest item while leaking padding into shorter items.
What shapes should you write on the whiteboard?
For a conventional multihead calculation, label scores as [batch, heads, L, S]. The final two dimensions are query and key length; they are not necessarily equal during cross-attention or cached decoding.
Official API fact: SDPA requires an attention mask broadcastable to the attention-score shape. MHA has its own documented two- or three-dimensional attention-mask layouts and a separate batch-by-key padding mask. See the respective SDPA and MHA references.
Original broadcast ledger: With scores [2, 4, 3, 5], a per-batch key-validity tensor [2, 5] should be expanded to [2, 1, 1, 5] for this explicit SDPA-style construction. Its singleton head and query dimensions repeat the same valid-key rule within each item.
Passing [2, 5] directly does not express that intended axis placement. Broadcasting aligns dimensions from the right, so the first dimension can align with the query axis rather than the batch axis, or fail to match it.
Assert semantic dimensions as well as rank. A tensor of shape [L, S] can still be transposed or built from the wrong positions. Shape checks establish compatibility; a small allowed-pair ledger establishes meaning.
Why does cached decoding need a different causal alignment?
Official implementation fact: A KV cache retains earlier key and value projections. New queries attend over the retained prefix plus current keys; the mask must cover that combined context. See Hugging Face's cache explanation.
Original exercise: Three keys at positions zero, one, and two are cached. Two new queries and keys represent absolute positions three and four. The score matrix is two by five, with this allowed relation:
key 0 1 2 3 4
query position 3 1 1 1 1 0
query position 4 1 1 1 1 1
A local two-by-five lower triangle would instead allow only the first key in the first row and the first two keys in the second. That matrix ignores the query offset and discards valid context.
Official API fact: PyTorch 2.14 SDPA documents an upper-left causal alignment for nonsquare is_causal=True masks. It also documents that supplying an explicit mask together with is_causal=True is an error. See the current SDPA parameter contract.
Engineering inference: For this cached fixture, construct the absolute-position relation explicitly and call with is_causal=False. Do not assume a flag that was correct for square prefill expresses the intended nonsquare decoding relation.
For a single new query at position four with keys zero through four, all five real keys are causally allowed. This does not authorize attending to unused slots in a preallocated cache. Validity and capacity are different quantities.
Position IDs and cache slots are also distinct. A padded batch, sliding window, or packed sequence can require a position mapping beyond past_length + local_index. Preserve the actual position and segment contract; never infer semantic positions solely from allocated tensor size.
Which small tests catch mask bugs?
Original test matrix: Compare full-sequence outputs with cached outputs for the same Q, K, and V, using deterministic settings and the correctly offset mask. They should agree within a justified numerical tolerance for corresponding query rows.
Then perturb forbidden future keys and values. Permitted earlier outputs should be invariant. Perturb padded keys and verify that real outputs remain unchanged. These tests can detect an inverted convention even when every tensor has a valid shape.
Use at least two batch items with different padding lengths, and include a multi-token cached chunk. A one-token test alone can miss a causal restriction between new tokens in the same chunk.
Check that attention weights on forbidden positions are zero in a transparent reference calculation, and that each nonempty allowed row sums to one. Compare with a small CPU implementation before interpreting differences from an optimized kernel.
Official operational fact: SDPA applies dropout according to its explicit dropout probability; evaluation mode alone does not rewrite that argument. Pass zero when the intended test is deterministic evaluation. See the API warning.
The article's local NumPy checks cover square attention, padded-key invariance, cached chunk equivalence, and a deliberately wrong local triangle. They verify this reference construction; they do not claim a full pretrained model or every PyTorch GPU backend was tested.
Five questions for implementation practice
These verified direct pages provide focused attention and cache exercises. They are not candidate reports establishing any employer's current round or interview frequency.
| Question | Mask follow-up |
|---|---|
| Implement Masked Multi-Head Self-Attention | Define causal and padded-key permission. |
| Implement correct attention masking | Check boolean conventions and axes. |
| Implement Multi-Head Attention with a KV Cache and Cached Grouped-Query Attention | Trace absolute query positions with cached keys. |
| Debug Transformer and Add KV Cache | Compare chunked and full-sequence outputs. |
| Implement Stable Sigmoid, Softmax, and Scaled Dot-Product Attention | Test masked softmax and empty rows. |
Continue with Implement Multi-Head Attention with a KV Cache and Cached Grouped-Query Attention. Draw the two-by-five relation before coding the cache update, then compare its result with full-sequence attention.