Attention-Gym: Triton-Based Sparse and Quantization Attention

Attention-Gym is a flexible and efficient framework built on Triton, designed to help researchers and developers rapidly implement, test, and validate innovative attention mechanisms. With support for sparse and quantized attention, it provides a powerful base environment for experimenting with new algorithms and optimizing existing ones.

Requirements

• python>=3.9 , torch>=2.3.0 , triton>=3.0.0, NVIDIA GPUs (Compute Capability 8.0+)
• Notice: FP8 dtype is only supported on NVIDIA GPUs (Compute Capability 9.0+)

Installation

pip install -e.

Kernels

Now Support:

• flash_attention2
• sliding_tile_attention
• sageattn_qk_int8_pv_fp16
• sageattn_qk_int8_pv_fp8
• sparge_sageattn_qk_int8_pv_fp16
• sparge_sageattn_qk_int8_pv_fp8

How to Use

To easy use:

import attention_gym
out = attention_gym.sageattn_qk_int8_pv_fp16_triton(q, k, v, tensor_layout="HND", is_causal=False)

• q, k, v are FP16/BF16 dtype with the shape (batch_size, head_num, seq_len, head_dim) using default tensor_layout="HND". For shape (batch_size, seq_len, head_num, head_dim), set tensor_layout="NHD".
• is_causal determines the use of a causal mask.

Kernel Tests

To run the tests:

pytest tests/test_sageattn_qk_int8_pv_fp16.py

Kernel Benchmark

To run the benchmarks:

python benchmarks/benchmark_sage1.py

End-to-end Performance And Accuracy

Here we compare the end-to-end performance and accuracy of the original algorithm author's CUDA implementation and the attention-gym triton implementation of each algorithm.

Acknowledgement

We learned the design and resued some code from the following projects: triton, FastVideo, SpargeAttn, SageAttention