LLM Systems
Invert boolean masks when migrating PyTorch attention APIs
A boolean mask passed unchanged from MultiheadAttention to scaled_dot_product_attention can reverse which keys participate. The safe conversion is small and easy to verify.
Evidence and testing
Review record: Tensor Almanac automated evidence gate (no independent human review). Example versions, recorded outputs and limitations appear below.
The boolean convention changes between APIs
A boolean mask is not automatically portable between two PyTorch attention APIs. For `nn.MultiheadAttention`, `True` in a `key_padding_mask` marks a key to ignore, and `True` in its boolean `attn_mask` marks a query-key position that is not allowed. For `torch.nn.functional.scaled_dot_product_attention`, a boolean `attn_mask` uses the opposite convention: `True` means the position participates. PyTorch calls out this inversion in the SDPA documentation and recommends logically negating a mask when migrating from MultiheadAttention. [1] [2]
That means a mask named `padding_mask` can look reasonable in a debugger while silently admitting padding and excluding real tokens after a direct API swap. The function name or tensor shape does not tell you which boolean convention is expected; check the receiving operation's contract.
Convert the meaning, then verify the weights
Suppose the third key is padding. In MultiheadAttention's convention, `[False, False, True]` means the first two keys are available and the last is blocked. SDPA expects a boolean keep-mask, so the equivalent input is `[True, True, False]`. A direct pass of the original values instead keeps only the padded key in this small example.
Run `python examples/attention_boolean_mask_conversion.py`. The NumPy fixture applies stable softmax to fixed scores and values after masking. With the mask converted, both interpretations produce the same weighted result, 12.69. Reusing the blocking mask as an SDPA keep-mask produces 1000.00 because the deliberately large sentinel value is the only one left. This is a semantics demonstration, not a measured PyTorch model output or performance claim.
Boolean inversion is the point of this example; dimensions still need to match the API being called. MultiheadAttention documents padding masks as batch-by-source and attention masks as target-by-source or batch-head-by-target-by-source. SDPA requires a mask broadcastable to its attention-weight shape. Float masks are additive score values, so do not invert a float mask as though it were boolean. [1] [2]
Executed with Python 3.12.14 and NumPy 2.3.5. The isolated example uses NumPy because PyTorch is not installed in this runtime; PyTorch API behavior was checked against the official documentation, not claimed as locally executed.
"""Model the boolean-mask conversion used when migrating attention APIs."""
import numpy as np
def attention_weights(scores, allowed):
values = np.asarray(scores, dtype=float)
keep = np.asarray(allowed)
if values.ndim != 1 or keep.shape != values.shape or keep.dtype != np.bool_:
raise ValueError("one-dimensional scores and matching boolean mask required")
if not keep.any():
raise ValueError("each query must retain at least one key")
masked = np.where(keep, values, -np.inf)
shifted = masked - np.max(masked)
weights = np.exp(shifted)
return weights / weights.sum()
def main():
scores = np.array([2.0, 1.0, 10.0])
values = np.array([10.0, 20.0, 1000.0])
mha_padding_mask = np.array([False, False, True])
mha_allowed = np.logical_not(mha_padding_mask)
sdpa_allowed = np.logical_not(mha_padding_mask)
wrong_sdpa_allowed = mha_padding_mask
mha_weights = attention_weights(scores, mha_allowed)
sdpa_weights = attention_weights(scores, sdpa_allowed)
wrong_weights = attention_weights(scores, wrong_sdpa_allowed)
mha_output = float(np.sum(mha_weights * values))
sdpa_output = float(np.sum(sdpa_weights * values))
wrong_output = float(np.sum(wrong_weights * values))
assert np.allclose(mha_weights, [0.73105858, 0.26894142, 0.0])
assert np.allclose(sdpa_weights, mha_weights)
assert np.allclose(wrong_weights, [0.0, 0.0, 1.0])
assert np.isclose(mha_output, sdpa_output)
assert np.isclose(wrong_output, 1000.0)
rejected = 0
try:
attention_weights(scores, np.array([False, False, False]))
except ValueError:
rejected += 1
assert rejected == 1
print("MultiheadAttention padding mask:", mha_padding_mask.tolist())
print("converted SDPA keep mask:", sdpa_allowed.tolist())
print(f"matching masked outputs: {mha_output:.2f} {sdpa_output:.2f}")
print(f"unconverted mask output: {wrong_output:.2f}")
print("all-blocked mask rejected:", rejected)
if __name__ == "__main__":
main()
Recorded output
MultiheadAttention padding mask: [False, False, True] converted SDPA keep mask: [True, True, False] matching masked outputs: 12.69 12.69 unconverted mask output: 1000.00 all-blocked mask rejected: 1
Keep a migration check beside the call site
When replacing MultiheadAttention with SDPA, make the conversion visible at the boundary: for a boolean mask whose `True` values mean blocked, pass its logical negation to SDPA. Then use a tiny fixture with one padded key and confirm its attention weight is zero while at least one valid key remains active. Such a check catches a convention change before a larger model or fused implementation makes the mistake harder to localize.
The executable fixture models one score row and one value vector. It does not exercise PyTorch tensor broadcasting, mask merging, causal flags, dropout, gradients, or an optimized kernel. It explicitly rejects an all-blocked row in its own normalization helper; it does not claim this is the behavior of every PyTorch backend. Follow each operation's current shape and mask documentation when adapting batched or multi-head code.