Batched transformers combine two masks: a causal mask (no peeking ahead) and a padding mask (ignore pad tokens in each sequence).
Implement attention_mask(lengths, T, causal=True) returning a boolean array of shape (B, T, T) where mask[b, i, j] is True if query i may attend to key j: key j must be a real token (j < lengths[b]) and, if causal, j <= i. Rows for padding queries can follow the same rule (they're discarded later).