Causal Mask as an Additive Tensor
def causal_mask(T, device):
mask = torch.zeros(T, T, device=device)
mask = mask.masked_fill(
torch.triu(torch.ones(T, T, device=device), diagonal=1).bool(),
float("-inf"),
)
return mask # (T, T), broadcasts to (B, H, T, T)
Why it matters: Additive masks (0.0 / -inf) compose — sum causal + padding and pass one tensor. Multiplicative masks don’t compose cleanly.