# TurboPairFormer: Fast and Stable Protein Folding Model Training with an Optimized Triangle Attention Kernel

arXiv: https://arxiv.org/abs/2610.05854 · HTML: https://mlsanity.com/paper/2610.05854/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 19 of its 33 references that are in the map.
- First posted on arXiv 5 Oct 2026.

## Builds on (most cited first, 19 of 19)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Self-attention Does Not Need O(n2) Memory](https://mlsanity.com/paper/2112.05682/) (2021)
- [FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling](https://mlsanity.com/paper/2603.05451/) (2026)
- [Nondeterminism and Instability in Neural Network Optimization](https://mlsanity.com/paper/2103.04514/) (2021)
- [DeepSpeed4Science Initiative: Enabling Large-Scale Scientific Discovery through Sophisticated AI System Technologies](https://mlsanity.com/paper/2310.04610/) (2023)
- [FastFold: Reducing AlphaFold Training Time from 11 Days to 67 Hours](https://mlsanity.com/paper/2203.00854/) (2022)
- [Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention](https://mlsanity.com/paper/2510.04212/) (2025)
- [FlashBias: Fast Computation of Attention with Bias](https://mlsanity.com/paper/2505.12044/) (2025)
- [Collage: Light-Weight Low-Precision Strategy for LLM Training](https://mlsanity.com/paper/2405.03637/) (2024)
- [ScaleFold: Reducing AlphaFold Initial Training Time to 10 Hours](https://mlsanity.com/paper/2404.11068/) (2024)
- [Attn-QAT: 4-Bit Attention With Quantization-Aware Training](https://mlsanity.com/paper/2603.00040/) (2026)
- [NeuralPLexer3: Accurate Biomolecular Complex Structure Prediction with Flow Models](https://mlsanity.com/paper/2412.10743/) (2024)
- [DASH: Deterministic Attention Scheduling for High-throughput Reproducible LLM Training](https://mlsanity.com/paper/2601.21824/) (2026)
- [MegaFold: Efficient Training of Next-Generation 3D Attention Protein Models on Cross-Platform GPUs](https://mlsanity.com/paper/2506.20686/) (2025)

## Abstract (arXiv)

Triangular attention is a core computation in AlphaFold3-style biomolecular models, with cubic cost in token count. Its shared pair bias adds a gradient reduction across attention slices to the usual reductions over queries and keys. The open-source backends we examine handle these reductions through repeated probability recomputation, floating-point atomics, or full score-gradient storage. Separately, computing the softmax backward correction from BF16-rounded forward outputs loses numerical precision. We present TurboPairFormer, a triangular attention implementation for NVIDIA Hopper GPUs that addresses both issues. Our key-tile-parallel backward algorithm recomputes each probability tile once for the query, key, value, and pair-bias gradients, using ordered partial reductions for deterministic accumulation without floating-point atomics or full score-gradient storage. Output-residual compensation retains a BF16 approximation of the output-rounding residual to compute the backward correction more accurately in FP32, without changing the BF16 output. With BF16 inputs at crop sizes 384, 640, and 768 and head dimensions 16 and 32, TurboPairFormer achieves the lowest mean query, key, and pair-bias gradient RMSE against an FP64 reference among the implementations compared in this paper. Residual compensation reduces these RMSE values by 28-47% in controlled ablations. All four gradients are bitwise identical across five repeated calls in all 600 input cases under fixed execution conditions. Integrated into OpenFold3 with our triangle multiplication kernels, TurboPairFormer achieves the lowest GPU computation time per optimizer step among the evaluated backend configurations on 16 H100 GPUs, with speedups of 1.73× over OpenFold3's Triton backend and 1.13× over cuEquivariance at crop size 768.