Your square causal mask is a clear full-sequence example. The KV-cache section is a useful place to show how that shape changes during decoding: one new query may attend to many cached keys, so the query and key sequence lengths are no longer equal.

When extending TinyCausalAttention with a cache, I would compare cached token-by-token outputs against a full-prefix forward pass in evaluation mode. That catches a mask aligned to the wrong absolute positions even when every tensor shape is accepted. It also connects the attention arithmetic directly to a serving correctness check.