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.

To reduce smem usage for SM75