javirub/flashattention-pytorch repository overview

⁠FlashAttention 2.7.4 + PyTorch 2.7.0 + CUDA 12.8 + cuDNN 9 (Runtime)

This image provides PyTorch 2.7.0 integrated with FlashAttention 2.7.4, compiled against CUDA 12.8 and cuDNN 9, optimized for high-performance training and inference on modern NVIDIA GPUs.

⁠📦 Image Details

• Base: pytorch/pytorch:2.7.0-cuda12.8-cudnn9-runtime
• PyTorch version: 2.7.0
• FlashAttention version: 2.7.4
• CUDA version: 12.8
• cuDNN version: 9
• Python version: 3.11

⁠🚀 How to Use

Pull the image: docker pull javirub/flashattention-pytorch:flashattn2.7.4-pytorch2.7.0-cuda12.8-cudnn9-runtime

Run the container with GPU support: docker run --gpus all -it javirub/flashattention-pytorch:flashattn2.7.4-pytorch2.7.0-cuda12.8-cudnn9-runtime

Inside the container, you can test FlashAttention installation: python3 >>> import flash_attn >>> print(flash_attn.version)

⁠⚡ Requirements

• NVIDIA GPU with CUDA 12.8+ support (Ampere, Ada or newer recommended)
• NVIDIA driver version 535+ (to support CUDA 12.8)
• Docker with --gpus support enabled

⁠📜 Notes

• This is a runtime image: it includes runtime libraries but no development headers.
• Suitable for both inference and training workloads that benefit from FlashAttention optimizations.

⁠📬 Contact

Created and maintained by javirub.

⁠👨🏻‍🎓 Citation

Uses https://github.com/Dao-AILab/flash-attention⁠

@inproceedings{dao2022flashattention,
  title={Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness},
  author={Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\'e}, Christopher},
  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
  year={2022}
}
@inproceedings{dao2023flashattention2,
  title={Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning},
  author={Dao, Tri},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2024}
}