GitLab

Update flash_attention.py

Recasting query and key after rotary_emb()

Phil Tillet suggests calling it "experimental".

Theoretically this might have lower numerical error since the scaling is in
fp32 instead of fp16 (not sure, I haven't thought too carefully about it).
However, in practice, the numerical errors seem about the same.

Slightly improve installation process

This speeds up the fwd by 1.5x.