Implement SageAttention for Efficient and Accurate 8-bit Attention

[ighoshsubho]

We should implement SageAttention, a novel method for efficient and accurate 8-bit attention computation, as described in the paper "SageAttention: Accurate 8-bit Attention for Plug-and-Play Inference Acceleration" (Zhang et al., 2024).

Key features to implement:

1. Smoothing matrix K to enhance accuracy
2. INT8 quantization for Q and K matrices
3. FP16 accumulator for P and V matrices
4. Adaptive quantization strategy
5. Fusion tricks for reduced overhead

Expected benefits:

• ~2.1x speedup compared to FlashAttention2
• ~2.7x speedup compared to xformers
• Negligible loss in end-to-end metrics across various models

Additional considerations:

• Ensure compatibility with existing model architectures
• Optimize for both RTX 4090 and 3090 GPUs
• Consider implementing for other GPU architectures in the future

References:

• Link to the SageAttention paper