[causal attn mask] replaced for loop over tensor with pytorch tensor ops

[dabeschte]

The original causal attention mask generation is very slow, especially when the tensor is created on the GPU, because it needs to make 1000s of calls.

I tried to compile it...which works and makes it fast too, but compilation unfortunately also takes a long time when using a long sequence length.

This implementation is ~20-70x faster depending on the sequence lengths and since it is re-created for every SDPA, this accumulates to multiple seconds per step for larger videos