Skip to content

Module 02 · Chapter 5

Causal and Padding Masks: Implementation and Tests

Construct causal and key-padding masks, combine via broadcast AND, and verify future-token leakage. Covers causal mask, padding mask, combined formula, and fully masked rows.

Intermediate15 min readUpdated Nov 1, 2026

Learning objectives

  • Construct the causal lower-triangular mask and the key padding mask for a given batch.
  • Combine the two masks via broadcast AND to produce the combined attention mask.
  • Run a leakage test that verifies attention weights are ~0 for all future-token positions.
  • Explain how key padding masks affect padded keys and queries, and how loss ignores padded targets.

Before you start

  • Comfort with the [B, T, d] tensor shape trace from the transformer-architecture lesson.
  • Basic understanding of scaled dot-product attention: softmax(QK^T / sqrt(d))V.
  • Familiarity with boolean tensor operations (AND, OR, NOT, broadcast) in PyTorch or NumPy.

Causal attention masks are the mechanism that prevents a token from “seeing” future tokens in a decoder-only, auto-regressive transformer. This lesson builds the two core masks (causal + key padding), combines them, and verifies correctness with a reproducible test. The lesson assumes the reader has already traced a small decoder-only model in the transformer-architecture lesson and knows the [B, T, d] shape convention.

A causal transformer applies a lower-triangular causal mask AND a key-padding mask (broadcast together) so that every attention weight at position t is either (a) a valid distribution over past-and-current real tokens, or (b) effectively zero for future or padded positions.

Everything else in this lesson makes that sentence concrete.

Lab: Run the companion CPU exercise at /labs/track1/01-attention/ (causal mask, padding mask, leakage test).

Use these dimensions throughout the lesson:

Symbol Meaning Value in this trace
B Batch size (sequences processed together) 2
T Sequence length (tokens per sequence) 5
h Number of attention heads 4
d_head Width per head (d / h) 8
d Model width (hidden dimension) 32

The numbers are deliberately small. A real model uses the same pattern with larger values; the pattern is what matters.

In a causal decoder-only model, position t may attend only to positions 0 through t. The mask is a lower-triangular matrix of shape [T, T] (or [1, 1, T, T] after broadcasting). We use the convention where True at (i, j) means keep / attend, and False means mask out (add -inf before softmax).

M_causal[i, j] = True if i >= j (position i can attend to position j)
M_causal[i, j] = False if i < j (position i cannot attend to position j)

In PyTorch:

causal = torch.tril(torch.ones(T, T, dtype=torch.bool))
# or equivalently:
causal = ~torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)

The mask is unsqueezed to [1, 1, T, T] so it can be AND-ed with the key padding mask and broadcast across batch and heads:

causal_4d = causal[None, None, :, :] # [1, 1, T, T]
j=0 j=1 j=2 j=3 j=4
i=0: T F F F F ← position 0 can attend only to 0
i=1: T T F F F ← position 1 can attend to 0, 1
i=2: T T T F F ← position 2 can attend to 0, 1, 2
i=3: T T T T F ← position 3 can attend to 0, 1, 2, 3
i=4: T T T T T ← position 4 can attend to 0, 1, 2, 3, 4

Key point: The causal mask is static — it depends only on T, not on the data. For any batch item and any head, the same lower-triangular pattern applies.

When sequences in a batch have different lengths, padding tokens must be excluded from attention. The key padding mask has shape [B, T_padded] (batch × sequence length), where True indicates a real token and False indicates padding:

attention_mask[b, t] = True if position t in batch b is a real token
attention_mask[b, t] = False if position t in batch b is padding

Convention note: This brief uses the Hugging Face DataCollatorForLanguageModeling(mlm=False) convention: True = real token, False = padding. This is the inverse of the attention_mask passed to the model during inference in some other kits. The lesson must state which convention is in force.

Given a list of sequence lengths and a pad_token_id, the mask is built by comparing each position to the pad_token_id:

# Example: B=2, T=5, real lengths = [3, 4]
# Token IDs: [[5, 12, 7, 0, 0], [3, 9, 1, 0, 0]] (0 = pad)
# attention_mask: [[True, True, True, False, False],
# [True, True, True, True, False]]

In PyTorch (Hugging Face convention):

# seq_lens: list of real lengths per batch item, padded to T
padding_mask = torch.zeros(B, T, dtype=torch.bool)
for b, L in enumerate(seq_lens):
padding_mask[b, :L] = True # True = real token

The mask is unsqueezed to [B, 1, 1, T] so it can be AND-ed with the causal mask and broadcast across heads and the key dimension:

key_pad_4d = key_pad[:, None, None, :] # [B, 1, 1, T]

| b\t | 0 | 1 | 2 | 3 | 4 | | —– | — | — | — | — | — | –––––––––––––––––––––– | | 0 | T | T | T | F | F | ← real tokens at positions 0,1,2; pad at 3,4 | | 1 | T | T | T | T | F | ← real tokens at positions 0,1,2,3; pad at 4 |

The combined mask that the attention computation applies is the AND of the causal mask and the key padding mask, broadcast to four dimensions:

combined_mask = causal_4d & key_pad_4d

where causal_4d has shape [1, 1, T, T] and key_pad_4d has shape [B, 1, 1, T]. The broadcast rule produces [B, 1, T, T]:

  • The 1 dimensions of causal_4d expand to B (batch) and 1 (head).
  • The 1 dimensions of key_pad_4d expand to B (batch) and T (key dimension, duplicated across the T positions of the causal matrix).
  • The AND is element-wise: combined_mask[b, h, t, j] = causal[t, j] AND key_pad[b, j] — the key padding mask selects which keys (position j) are real, independent of which query position t is attending.

A common error is to OR the two masks: combined = causal | key_pad_expanded. This would let padding positions attend causally — a silent correctness bug. The AND ensures that:

  • A real query at position t can attend to a real key at position j only if j <= t (causal) and key j is not padding (key_pad).
  • A real query at position t cannot attend to a future key (causal), even if that key is real.
  • A query at a padding position has its entire row masked (AND with False yields False everywhere).

3.2 Worked example (B = 2, T = 5, from §2.2)

Section titled “3.2 Worked example (B = 2, T = 5, from §2.2)”

Causal mask (lower-triangular, True = attend):

j=0 j=1 j=2 j=3 j=4
i=0: T T T T T
i=1: F T T T T
i=2: F F T T T
i=3: F F F T T
i=4: F F F F T

Key padding mask (from §2.2):

b=0: T T T F F
b=1: T T T T F

Combined mask (broadcast AND → [B, 1, T, T]):

b=0:
j=0 j=1 j=2 j=3 j=4
i=0: T T T F F ← causal T & pad T,T,T,F,F → T,T,T,F,F
i=1: F T T F F ← causal F&T,T,T,T,T & T,T,T,F,F → F,T,T,F,F
i=2: F F T F F ← causal F,F,T,T,T & T,T,T,F,F → F,F,T,F,F
i=3: F F F F F ← causal F,F,F,T,T & T,T,T,F,F → F,F,F,F,F
i=4: F F F F F ← causal F,F,F,F,T & T,T,T,F,F → F,F,F,F,F
b=1:
j=0 j=1 j=2 j=3 j=4
i=0: T T T T F ← causal T & pad T,T,T,T,F → T,T,T,T,F
i=1: F T T T F ← causal F&T,T,T,T,T & T,T,T,T,F → F,T,T,T,F
i=2: F F T T F ← causal F,F,T,T,T & T,T,T,T,F → F,F,T,T,F
i=3: F F F T F ← causal F,F,F,T,T & T,T,T,T,F → F,F,F,T,F
i=4: F F F F F ← causal F,F,F,F,T & T,T,T,T,F → F,F,F,F,F

Engineering judgment: Study the b=0, i=3 row: the causal mask alone would let position 3 attend to positions 3 and 4. But position 4 is padding (key_pad[b=0, 4] = False), so the combined mask zeroes it out. Similarly, b=0, i=1: causal alone would let position 1 attend to position 0 (which is F in causal, so it’s already masked). The padding mask further restricts the allowed zone.

A robust way to verify mask correctness is to run a forward pass with a known input and check that the attention weight matrix has zeros (or near-zeros) in the future-token positions.

  1. Set up a minimal batch: B = 1, T = 4, with two padded positions at the end. Use d = 16, h = 2, d_h = 8.
  2. Compute Q, K, V randomly (shaped [B, h, T, d_h] where d_h = d / h).
  3. Build the combined mask per §3 (causal + key padding, broadcast AND).
  4. Compute attention scores: attn_scores = Q @ K.transpose(-2, -1) / sqrt(d_h).
  5. Apply the combined mask: attn_scores = attn_scores.masked_fill(~combined_mask, float('-inf')).
  6. Apply softmax: attn_weights = torch.softmax(attn_scores, dim=-1).
  7. Verify:
    • attn_weights[b, h, t, j] ≈ 0 for all j > t (future positions are masked).
    • attn_weights[b, h, t, j] is finite and > 0 for j <= t and position j is not padding.
    • Padded keys (j is pad) get weight ~0 from ALL queries.
    • Padded queries (t is pad) attend to real keys normally — loss ignores them via labels=-100.

4.2 Why softmax of an all--inf row yields NaN (plain softmax) vs. SDPA

Section titled “4.2 Why softmax of an all--inf row yields NaN (plain softmax) vs. SDPA”

Before showing the test, one important subtlety: if a query row is entirely masked (every key position is -inf), plain torch.softmax computes exp(-inf) = 0 for every entry, then divides by sum = 0, producing 0/0 = NaN.

This happens only when every valid key for a query is itself padded. For example, if position 0 is padded (key_pad[0] = False), position 0 can only attend to itself via causal masking — but itself is padded — so the entire row is -inf, producing NaN with plain softmax.

PyTorch SDPA behaves differently. torch.nn.functional.scaled_dot_product_attention (tested on PyTorch 2.15.0) does not return NaN for a fully masked row — it returns an all-zero output row for that query. So if you use SDPA instead of manual softmax, a fully masked row yields zeros, not NaN. Verify this yourself on your target PyTorch version before relying on NaN-handling code.

For ordinary padded queries at positions 2,3 (which CAN attend to real keys 0,1), the row is NOT all -inf — positions 0 and 1 are real keys — so softmax produces valid weights, not NaN. Production code handles NaN rows (plain softmax) via:

  1. Skip computing attention for padded queries entirely (varlen kernels do this).
  2. Replace NaN rows with zeros after softmax (torch.where(torch.isnan(attn_weights), 0, attn_weights)).
  3. Ensure padded queries never reach the attention computation (forward pass short-circuits them).

It is critical to distinguish the two:

Padded key (position j is pad) Padded query (position t is pad)
Mask effect combined[b,:,t,j] = False for ALL queries t — column j is zeroed key_pad[b,t] = False — row t is NOT itself masked by the padding mask
Attention weight All queries give weight ~0 to padded key j Padded query t attends to real keys normally
Output contribution Padded key contributes nothing to any context vector Padded query produces valid output; loss ignores via labels=-100

The combined mask causal_4d & key_pad_4d correctly zeroes padded keys as columns. Padded queries still attend to real keys — their output is valid but discarded by the loss function via labels=-100.

import torch
# ---- Setup ----
B, T, d, h = 1, 4, 16, 2 # batch, size, width, heads
d_h = d // h # 8 — score scaling uses head dimension, not model width
torch.manual_seed(42)
Q = torch.randn(B, h, T, d_h) # [1, 2, 4, 8] — [B, h, T, d_head]
K = torch.randn(B, h, T, d_h)
V = torch.randn(B, h, T, d_h)
# ---- Masks ----
# Causal: lower triangular, True = attend
causal = torch.tril(torch.ones(T, T, dtype=torch.bool)) # [4, 4]
# Key padding mask: True = real, False = pad
# Example: T=4, real tokens at positions 0,1; pad at positions 2,3
key_pad = torch.tensor([[1, 1, 0, 0]], dtype=torch.bool) # [1, 4]
# Expand for broadcast
causal_4d = causal[None, None, :, :] # [1, 1, 4, 4]
key_pad_4d = key_pad[:, None, None, :] # [1, 1, 1, 4]
# Combined: broadcast AND → [1, 1, 4, 4]
combined = causal_4d & key_pad_4d
# ---- Attention ----
# Score scaling by sqrt(d_head), per PyTorch docs for scaled_dot_product_attention
attn_scores = Q @ K.transpose(-2, -1) / (d_h ** 0.5) # [1, 2, 4, 4]
attn_scores = attn_scores.masked_fill(~combined, float('-inf'))
attn_weights = torch.softmax(attn_scores, dim=-1)
# ---- Verify ----
print("=== Attention weights (B=1, h=2, T=4) ===")
print(attn_weights.squeeze(0)) # shape [2, 4, 4]
print("\n=== Verification ===")
# Check 1: Future positions (j > t) should be ~0 for ALL queries and heads
for t in range(T):
for j in range(T):
if j > t:
for h_idx in range(h):
w = attn_weights[0, h_idx, t, j].item()
assert w < 0.01, f"FAIL: head {h_idx}, position {t} attends to future {j} with weight {w}"
print("PASS: All future positions (j > t) masked for every query and head.")
# Check 2: Padded keys (j=2,3) have weight ~0 for ALL queries and heads
for h_idx in range(h):
for t in range(T):
for j in [2, 3]:
w = attn_weights[0, h_idx, t, j].item()
assert w < 0.01, f"FAIL: head {h_idx}, query {t} gives weight {w} to padded key {j}"
print("PASS: Padded keys (j=2,3) have zero attention weight for all queries and heads.")
# Check 3: Real non-future keys have valid weights (> 0)
for h_idx in range(h):
for t in range(T):
for j in range(t + 1): # j <= t
if j < 2: # only real keys
w = attn_weights[0, h_idx, t, j].item()
assert w > 0.01, f"FAIL: head {h_idx}, query {t}, real key {j} has weight {w}"
print("PASS: Real non-future keys have valid attention weights.")
# Check 4: Padded queries (t=2,3) attend to real keys — they DO NOT produce NaN
# because real keys 0,1 are available. Loss ignores them via labels=-100.
for h_idx in range(h):
for t in [2, 3]: # padded queries
row = attn_weights[0, h_idx, t, :]
# Padded query attends to real keys 0,1 (not future, not padded)
assert row[0] > 0.01, f"FAIL: head {h_idx}, padded query {t} should attend to real key 0"
assert row[1] > 0.01, f"FAIL: head {h_idx}, padded query {t} should attend to real key 1"
# Padded keys get zero weight
assert row[2] < 0.01 and row[3] < 0.01
print("PASS: Padded queries attend to real keys normally; padded keys masked; loss ignores via labels=-100.")
print("\n=== All checks passed ===")

Expected console output (invariants — exact values depend on Q/K/V seed):

=== Attention weights (B=1, h=2, T=4) ===
# head 0:
tensor([[1.0000, 0.0000, 0.0000, 0.0000], # t=0: attends to j=0 only (causal + real key)
[0.5000, 0.5000, 0.0000, 0.0000], # t=1: attends to j=0,1 (real, not future)
[0.6000, 0.4000, 0.0000, 0.0000], # t=2: padded query → attends to real keys 0,1
[0.7000, 0.3000, 0.0000, 0.0000]]) # t=3: padded query → attends to real keys 0,1
# head 1: similar pattern
=== Verification ===
PASS: All future positions (j > t) masked for every query and head.
PASS: Padded keys (j=2,3) have zero attention weight for all queries and heads.
PASS: Real non-future keys have valid attention weights.
PASS: Padded queries attend to real keys normally; padded keys masked; loss ignores via labels=-100.
=== All checks passed ===
  • Missing causal mask: future positions have non-zero weights → FAIL.
  • Missing padding mask: real future tokens attend with non-zero weights → FAIL.
  • OR instead of AND: padding positions attend causally → FAIL (e.g., b=0, i=3 attends to j=4 which is padding but gets a non-zero weight).
  • Wrong Q/K/V shape: Q @ K.transpose(-2, -1) with [B,T,h,d_h] yields [B,T,h,T] instead of [B,h,T,T] → shape mismatch or wrong attention → FAIL.
  • Wrong score scaling: dividing by d ** 0.5 (model width) instead of d_h ** 0.5 (head dimension) → over/under-scaled logits → softmax too sharp/flat → FAIL.
  • Wrong broadcast shape: TypeError: The size of tensor a (T) must match the size of tensor b (T_padded) → FAIL (crash, not silent).
  • Inverted padding mask convention: True = pad, False = real → real keys get zero weights, future tokens attend → FAIL (opposite of expected).

4.6 Engineering judgment (Staff/Principal)

Section titled “4.6 Engineering judgment (Staff/Principal)”

The future-token leakage test is the single most important verification step when implementing or debugging causal attention masks. It is inexpensive (one forward pass) and catches all three classes of mask errors: missing causal, missing padding, and wrong combination (AND vs. OR). Include this test in every PR that changes mask construction, and make it a gate in the CI pipeline.

Aspect Causal mask Key padding mask
Purpose Prevent future-token leakage (auto-regressive) Exclude pad tokens from attention
Shape (before broadcast) [T, T] → [1, 1, T, T] [B, T] → [B, 1, 1, T]
Value semantics True at (i, j) if i >= j (can attend) True at (b, t) if position t is real
Data-dependent No (static, same for every batch) Yes (varies per batch)
Broadcast dimensions [1, 1, T, T] → [B, 1, T, T] [B, 1, 1, T] → [B, 1, T, T]
Combined formula combined = causal_4d & key_pad_4d N/A
Typical use Both training and inference Only training with variable-length batches
During inference (T=1) Always 1×1 lower-trivial Not needed (no padding at inference)
Common error Forgetting to invert convention (OR vs. AND) Using the wrong True/False semantics

NaN occurs only when a query position has no valid keys at all — i.e., every key position j <= t (causal constraint) is itself padded. In that case, plain torch.softmax computes exp(-inf) = 0 for every entry, then divides by sum = 0, producing 0/0 = NaN. This happens when the very first position is padded (position 0 can only attend to itself, which is padded) — not for ordinary padded queries that can still attend to earlier real keys.

PyTorch SDPA behaves differently. torch.nn.functional.scaled_dot_product_attention (tested on PyTorch 2.15.0) returns an all-zero output row for a fully masked query, not NaN. So the “NaN” behavior above applies to manual torch.softmax; if you use SDPA, fully masked rows yield zeros instead.

  • Ordinary padded query (e.g., position 2 when positions 0,1 are real): attends to real keys 0,1 normally. Valid output. Loss ignores via labels=-100. No NaN.
  • Fully masked row (e.g., position 0 when position 0 itself is padded): position 0 can only attend to position 0, which is padded → entire row -inf → softmax → NaN.

Production code handles NaN rows via:

  1. Skip computing attention for padded queries entirely (varlen kernels do this).
  2. Replace NaN rows with zeros after softmax (torch.where(torch.isnan(attn_weights), 0, attn_weights)).
  3. Ensure padded queries never reach the attention computation (forward pass short-circuits them).

The loss function (with labels = -100 for pad positions) ignores padded positions, but only if NaN has been replaced or skipped — NaN * 0 in cross-entropy is still NaN.

Staff/Principal engineering judgment: Fully masked rows are unsafe by default — with plain softmax they produce NaN, with SDPA they produce zeros. A common production bug is assuming padded-query rows yield zero attention weights (they don’t with plain softmax; they yield NaN when fully masked). The fix is to zero-replace NaN rows after softmax or to skip padded queries before attention. If you see NaN in attention weights or loss, check whether fully-masked rows were handled. The mask construction itself (AND of causal + key_pad) is correct; the NaN handling is a separate production concern.

Problem: Given the following batch configuration, construct all masks and verify the future-token leakage test.

  • B = 2, T = 5, d = 16, h = 4, d_head = 4
  • Sequence lengths: [3, 4] (batch item 0 has 3 real tokens then 2 pads; batch item 1 has 4 real tokens then 1 pad)
  • Token IDs (pad token = 0): [[5, 12, 7, 0, 0], [3, 9, 1, 0, 0]]

Steps:

  1. Build the key padding mask (attention_mask) for both batch items, using the Hugging Face convention (True = real token).
  2. Build the causal mask (causal) for T = 5, then broadcast to [1, 1, 5, 5].
  3. Broadcast the key padding mask to [B, 1, 1, T], then AND with the causal mask to produce the combined mask (combined) of shape [B, 1, 5, 5].
  4. Write the PyTorch smoke-test code (as in §4.2) that verifies future-position weights are ~0, padded keys have ~0 weight from all queries, padded queries attend to real keys normally, and non-future real positions > 0.
  5. What-if: If the combined mask were computed with OR instead of AND, which entries in the combined mask would differ, and what failure mode would that produce? Answer: The OR would let padding positions attend causally. For b=0, i=3, the OR would keep True at j=4 (since causal[3,4]=True and pad[b=0,4]=False, OR gives True). The test would FAIL because position 3 would have a non-zero attention weight at future position 4, even though position 4 is padding.

Exercise answers (instructor):

  1. Key padding mask:

    • b=0: [True, True, True, False, False] (real at 0,1,2; pad at 3,4)
    • b=1: [True, True, True, True, False] (real at 0,1,2,3; pad at 4)
  2. Causal mask (T=5, lower-triangular, True = attend):

    j=0 j=1 j=2 j=3 j=4
    i=0: T T T T T
    i=1: F T T T T
    i=2: F F T T T
    i=3: F F F T T
    i=4: F F F F T
  3. Combined mask (AND broadcast → [2, 1, 5, 5]):

    • b=0: same pattern as the worked example in §3.2 (from the causal & pad[T,T,T,F,F] row).
    • b=1: causal AND pad[T,T,T,T,F] → each row AND-ed with the pad pattern.
  4. Smoke-test code: See §4.2. Expected output: future positions ~0, padded keys ~0 weight from all queries, padded queries attend to real keys normally, non-future real positions > 0.

  5. What-if: OR instead of AND. For b=0, i=3: causal[3,4]=True, pad[b=0,4]=False → OR gives True. The test would FAIL because position 3 would have a non-zero attention weight at future position 4. The failure mode: the model “sees” a future token (position 4 is padding, but the attention weight is non-zero), which constitutes future-token leakage.

    Post-fix instructions: After correcting OR back to AND, re-run the smoke test from §4.2 and verify all four checks pass: (1) future positions have weight ~0 for every query and head; (2) padded keys (j=3,4 for b=0; j=4 for b=1) have weight ~0 from all queries; (3) real non-future keys have weight > 0; (4) padded queries attend to real keys normally. If any check fails, the combined mask still has an indexing or broadcast error — re-verify that key_pad indexes the key dimension j, not the query dimension t.

Causal + Padding Mask Construction

Causal and padding mask constructionThree panels. Panel A: causal lower-triangular mask for T=5, True=attend, strictly lower-triangular. Panel B: key padding mask for B=2, T=5: True=real token, False=pad. Panel C: combined mask (causal AND key_pad broadcast) and attention weight matrix from the leakage test: future positions ~0, padded rows all zero, non-future real positions have valid weights.Panel A: Causal mask (T=5)M[i,j] = True if i >= j, else FalseTF0123401234Panel B: Key padding mask (B=2, T=5)True = real token, False = pad (HF DataCollator convention)TFb=0TFb=1Panel C: Combined mask + leakage testCombined: causal AND key_pad broadcast; leakage test: future ~0, padded rows zeroTTTFFFTTTFFFTTFFFFFFFFFFFTTTTFTTTTFTTTTFTTTTFFTTFLeakage test results:Future positions (j > t) → attention weight ≈ 0Padded keys (j=3,4) → attention weight ≈ 0 for all queriesNon-future real positions → attention weight > 0
Full pipeline: causal mask → key padding mask → combined mask → leakage test
Read this diagram as text

Text equivalent: Panel A shows the causal lower-triangular mask for T=5: position i can attend to positions 0..i. Panel B shows the key padding mask for B=2, T=5: True=real token, False=pad. Positions 3,4 are pad for b=0; positions 4 is pad for b=1. Panel C shows the combined mask: causal [1,1,5,5] AND key_pad [2,1,1,5] → [2,1,5,5], with the broadcast AND rule making each batch's mask independent. Panel D shows the leakage-test attention weight matrix: future positions (upper-right triangle) are ~0; padded keys have ~0 weight from all queries; padded queries attend to real keys normally; non-future real positions have valid weights. The engineering judgment: include the future-token leakage test as a CI gate for any mask change.

Text equivalent: The diagram has four panels. Panel A illustrates the causal lower-triangular mask for T=5: position i can attend to positions 0 through i. Panel B shows the key padding mask for B=2, T=5: True = real token, False = pad. For b=0, positions 3 and 4 are padding; for b=1, position 4 is padding. Panel C combines the two: the causal mask [1,1,5,5] is AND-ed with the key padding mask [2,1,1,5] via broadcasting, producing the combined mask [2,1,5,5]. The broadcast AND rule means each batch item’s mask is independent — b=0’s mask uses b=0’s padding pattern, and b=1’s mask uses b=1’s. Panel D shows the attention weight matrix from the leakage test: the upper-right triangle (future positions) is ~0; the bottom row of each batch (padded positions) is all zero; non-future real positions have valid ( > 0) weights. The engineering judgment summarized below should guide the decision.

Question Correct approach Why it matters
Causal vs. OR combination combined = causal & key_pad_expanded (AND) OR would let padding positions attend causally, a silent correctness bug that produces wrong logits.
Padding mask convention True = real token, False = pad (HF DataCollatorForLanguageModeling(mlm=False)) The opposite convention (True = pad) appears in some other kits; using the wrong convention flips all mask semantics and breaks the test.
Fully masked rows Plain softmax → NaN; SDPA → all-zero row. Zero-replace NaN or skip padded queries. NaN with plain softmax means the row was fully masked; verify mask construction and NaN handling before training.
Future-token leakage test Include as a CI gate for any mask change One forward pass catches missing causal, missing padding, and wrong AND/OR combination — the single most important verification step.
Causal mask during inference Always applied, with T = 1 at each auto-regressive step The causal mask is not only for training; it is the default mask for generation.

Staff/Principal judgment: For any decoder-only model in the 100M–1B parameter range, always construct the combined mask as causal_4d & key_pad_4d, use the Hugging Face DataCollatorForLanguageModeling(mlm=False) convention for the key padding mask (True = real token), and include the future-token leakage test as a mandatory CI gate. The FLOP overhead of the combined mask is negligible; the correctness risk of getting it wrong is high (silent logit errors that are hard to debug). Document the mask-combining formula and the test procedure in the model’s ADR (architecture decision record).