GitLab

To avoid import error if one doesn't have rotary_emb installed

This is faster since we only need to do atomic adds on dq, instead of atomic
adds on both dk and dv.

WIP because there seems to be some race conditions for head dimensions other
than 16, 32, 64, 128.

They don't have to have the same block size, number of threads, etc.

According to the `setup.py` file, only dependencies are torch and einops. But the `bert_padding.py` file requires `numpy` only to multiply the elements of a `torch.Size` object. This change aims at allowing the use of FlashAttention without numpy.

Recasting query and key after rotary_emb()