Neural Networks & LLM Components · Medium

PyTorch Causal Mask

Build a batched boolean decoder self-attention mask from causal order and key padding.

Task

Implement build_causal_attention_mask(key_padding_mask: torch.Tensor) -> torch.Tensor. Build a batched boolean causal self-attention mask from a key padding mask.

Requirements

key_padding_mask has shape (B, T), where True means the key token is real and False means padding.
Return a boolean tensor with shape (B, T, T).
mask[b, q, k] must be True only when k <= q and key_padding_mask[b, k] is True.
Future key positions must be False.
Padding key positions must be False for every query.
Keep the mask on the same device as key_padding_mask.

Example

build_causal_attention_mask(torch.tensor([[True, True, False]]))
# tensor([[[ True, False, False],
#          [ True,  True, False],
#          [ True,  True, False]]])