This repository contains an efficient implementation of triangular attention using Triton, specifically designed for protein structure prediction models like AlphaFold2.
The implementation is based on the Flash Attention 2.0 algorithm but modified to handle the unique requirements of triangular attention: 5d tensors operating over pairwise representations (B, H, L, L, D), and a bias term over the attention weights, through which we wish to propagate gradients.
- bf16 and fp32 support
- attention weight dropout
- Supporting for incoming edges (currently must manually permute and call kernel)
- TMA support
- Query (Q):
(Batch, Heads, L1, L2, Dims) - Key (K):
(Batch, Heads, L1, L2, Dims) - Value (V):
(Batch, Heads, L1, L2, Dims) - Bias (B):
(Batch, Heads, L1, L2)
- Additional Bias Term: Unlike standard attention which only uses QK^T similarity, this implementation includes an additional bias term that modulates the attention scores.
- 5D Tensors: Works with 5D tensors (B, H, L1, L2, D) instead of the standard 4D tensors used in transformer attention
Benchmarked on an NVIDIA L4 GPU with the following parameters:
- Batch Size: 4
- Heads: 8
- Head Dimension: 32
See example_usage.py for a simple example of how to use the implementation.
This implementation is an adaptation of the Triton tutorial implementation of Flash Attention 2.0, which can be found here. The attention mechanism is based on the AlphaFold2 paper. This repository is intended to be an open source reproduction of the flash triangular attention mechanism described in the NeuralPLexer 3 paper
If you use this codebase, or otherwise found it useful, please cite:
@software{triangular_attention_triton,
author = {Bryan Wu},
title = {Triton Implementation of Triangular Attention},
url = {https://github.com/bryanzwu/triangular-attention-triton},
year = {2025},
note = {An efficient implementation of triangular attention using Triton}
}This project is licensed under the MIT License. See the LICENSE file for details.