We should evaluate using Pytorch's built-in Flash Attention operator which also is compatible with their Nested Tensor which we could trivially wrap our JaggedTensor with. We should use Flash Attention when hardware and data types allow.