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.