Add Sage Attention for Lumina

[rockerBOO]

Requires https://github.com/thu-ml/SageAttention

• Sage Attention 2: 3.8s/it
• Flash Attention 2: 4.33s/it

Actual numbers may be different as not a full run to estimate with.

[sdbds]

sage attention cant use in training.
I think add SpargeAttn
can get more speed.

[rockerBOO]

Hmm good point but it does seem to be "training" the LoRA. The weights are present and seem to have actual weights. I will look a little deeper to see what is actually going on. Maybe it's just keeping the initialized weights or something.

thu-ml/SageAttention#60

SpargeAttn might require a significant "autotune" to produce a sparse attention model. thu-ml/SpargeAttn#6 I was initially exploring adding SpargeAttn but found that issue.

[rockerBOO]

Ok it seems it isn't adjusting the attention parameters but the others would be updated (like MLP/feed forward). I think keeping SageAttention could be fine but noting it is not for training? Possibly some work-arounds for the backwards that were applied for flash attention might make it work in the future.