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

Triangular attention is a core computation in AlphaFold3-style biomolecular models, with computational cost that grows cubically with token count. FlashAttention reduces attention memory traffic through IO-aware tiling and backward recomputation. However, the shared pair bias in triangular attention 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 incurs a loss of numerical precision. We present TurboPairFormer, a triangular attention implementation for NVIDIA Hopper GPUs that addresses these two issues. Our key tile parallel backward algorithm recomputes each probability tile once and reuses it for the query, key, value, and pair-bias gradients, using ordered partial reductions to achieve deterministic accumulation without floating point atomics or full score gradient storage. Our output residual compensation retains a BF16 approximation of the output rounding residual and uses it to compute the backward correction more accurately in FP32, without changing the BF16 output. With BF16 inputs at crop sizes 384, 640, 768 and head dimensions 16, 32, TurboPairFormer achieves the lowest mean query, key, and pair-bias gradient RMSE against an FP64 reference among eight evaluated implementations; residual compensation mitigates the precision loss from output rounding, reducing 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 together with our triangle multiplication kernels, TurboPairFormer achieves the lowest GPU computation time per optimizer step among eight evaluated backend configurations on 16 H100 GPUs, with speedups of $1.73\times$ over OpenFold3's Triton backend and $1.13\times$ over cuEquivariance at crop size 768.