发表机构
Stevens Institute of Technology; Lambda; Open Molecular Software Foundation（史蒂文斯理工学院; Lambda; 开放分子软件基金会）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
TurboPairFormer提出优化的三角注意力核，解决AlphaFold3式模型训练中梯度归约和BF16精度问题，在Hopper GPU上实现确定性累积和残差补偿，显著提升训练速度与梯度精度。
AI 中文摘要
三角注意力是AlphaFold3式生物分子模型中的核心计算，其计算复杂度随令牌数量呈三次方增长。其共享的成对偏置在通常的查询和键归约之上，增加了跨注意力切片的梯度归约。我们所考察的开源后端通过重复的概率重计算、浮点原子操作或完整的分数梯度存储来处理这些归约。此外，从BF16舍入的前向输出计算softmax反向修正会损失数值精度。我们提出了TurboPairFormer，一种面向NVIDIA Hopper GPU的三角注意力实现，解决了这两个问题。我们的键瓦片并行反向算法为查询、键、值和成对偏置梯度各重计算一次概率瓦片，使用有序部分归约实现确定性累积，无需浮点原子操作或完整的分数梯度存储。输出残差补偿保留了输出舍入残差的BF16近似，以在FP32中更精确地计算反向修正，同时不改变BF16输出。在裁剪尺寸为384、640和768、头维度为16和32的BF16输入下，TurboPairFormer在本文比较的实现中，相对于FP64参考实现了最低的平均查询、键和成对偏置梯度均方根误差。在受控消融实验中，残差补偿将这些均方根误差值降低了28-47%。在固定执行条件下，所有600个输入案例中，五次重复调用中四个梯度均逐位相同。集成到OpenFold3中并结合我们的三角乘法核，TurboPairFormer在16块H100 GPU上实现了评估的后端配置中每个优化器步骤最低的GPU计算时间，相比OpenFold3的Triton后端加速1.73倍，相比cuEquivariance在裁剪尺寸768时加速1.13倍。
英文摘要
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\times$ over OpenFold3's Triton backend and $1.13\times$ over cuEquivariance at crop size 768.
Comments33 pages, including appendices. Software: https://pypi.org/project/turbopairformer
