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

Yide Ran; Chelsea Lowman; Jan Domański; David Hartmann; Jenke Scheen; Jennifer Wei; Chuan Li; Jianwen Xie

cs.LGOct 5, 2026

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

Authors: Yide Ran, Chelsea Lowman, Jan Domański, David Hartmann, Jenke Scheen, Jennifer Wei, Chuan Li, Jianwen Xie, +1 more

Organizations: Stevens Institute of Technology · Lambda · Open Molecular Software Foundation

Abstract

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.

Figures & tables

Figure 1: Triangular attention and backward design comparison. (a) Triangular attention consists of

N

slices sharing one pair bias. Its backward pass involves three reduction axes:

dQ

over keys,

dK

and

dV

over queries, and

d\,\mathrm{bias}

across outer slices. All four gradients share the same intermediates (

P

and

dS

). (b) Prior backends either recompute probabilities 3

\times

, use non-deterministic atomics, or store the full score gradient

dS

, whose size is cubic in

N

. TurboPairFormer evaluates

P

and

dS

only once: it keeps

dK

and

dV

in local registers, writes

dQ

and

d\,\mathrm{bias}

into independent partial buffers, and reduces them in a deterministic FP32 order.

Figure 2: Work partitioning of the backward kernel. (a) A thread block owns one key tile in each of the four slices of an outer group and visits every query tile. (b)

dK

and

dV

finish inside the block;

dQ

and

d\,\mathrm{bias}

leave as partials with one writer each, which separate kernels add in a fixed order.

Figure 3: OpenFold3 training performance on 16

\times

H100 GPUs. Mean GPU computation time per optimizer step on two nodes with eight H100 GPUs each and 16 samples per step. Statistics cover steps 21–200; data loading, host-to-device transfer, and EMA updates are excluded. Lower is better. TurboPairFormer has the lowest measured mean at all three crop sizes.

Figure 4: Attention backward accuracy against FP64. Rows show

D=16

and

D=32

; columns show

dQ

dK

dV

, and

d\,\mathrm{bias}

. Points and error bars report mean absolute RMSE and sample standard deviation over ten seeds. Logarithmic axes show RMSE in units of

10^{-3}

. Implementation variants and additional results are described in Appendix G.3 .

Figure 5: Preview2 validation after fine-tuning on 16 H100 GPUs. Points show means and sample SDs over five evaluation seeds on 85 kinase validation complexes after 1,000 global training steps, with step 0 as a reference. Structure metrics average top-1 and best-of-five values; model selection combines lDDT and RASA means (Appendix G.5 ). Higher is better except for RMSD.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Notation	Meaning
$B$ , $N$ , $N_{o}$	Examples per GPU, tokens per crop, and number of fixed-token attention slices. The label $o$ in $N_{o}$ denotes the outer axis; $N_{o}=N$ .
$H$ , $D$ , $C_{z}$	Heads, features per head, and features per pair; the label $z$ refers to the pair tensor $Z$ .
$o,q,k$ ; $i,j$	Token indices: the fixed token (outer index), and the query and key positions within its slice. The triangle example uses $i=o$ and $j=q$ . A subscript $o$ alone selects one fixed-token slice.
$d$ ; $t,\ell,c$	Feature coordinate; key tile, outer group, and forward 16-key chunk indices. There are four chunks per forward tile. In the memory analysis, $t$ denotes an execution point.
$Z$	Pair features of shape $[B,N,N,C_{z}]$ . $Z_{ij:}$ selects one pair, and $Z_{o,:}$ (also written $Z_{o,:,:}$ ) the slice of a fixed token $o$ ; numeric subscripts such as $Z_{2,4,:}$ select individual pairs.
$Q,K,V,O$	Query, key, value, and attention output tensors, each of shape $[B,N_{o},H,N,D]$ ; for one example, slice, and head, each is an $N\times D$ matrix.

Appendix

Table 1: Symbols used in the method and appendix. Tensor shapes are logical; batch and head indices are omitted where stated.

Figure 6: Worked example of triangular attention. We follow one attention score through a crop of seven tokens with one head; numbers in subscripts are token positions. (1)

Z

stores one feature vector for every pair of tokens, and fixing token

o=2

takes row 2 of

Z

as one slice. (2) Each pair in this row is projected into a query, a key, and a value; the pair

(2,4)

gives the query

Q_{4,:}

, and the pair

(2,6)

gives the key

K_{6,:}

and the value

V_{6,:}

. (3) Their score also adds a bias taken from the third pair,

(4,6)

, and every slice uses this same bias for its positions 4 and 6. (4) The seven scores of query 4 go through a softmax, and the resulting weights average the seven values into the output that updates the pair

(2,4)

. (5) In the backward pass, this one score enters three different sums: over keys for

dQ

, over queries for

dK

and

dV

, and over slices for

d\,\mathrm{bias}

. Other queries and slices work the same way; masks and the normalization before the projections are omitted.

Figure 7: Backward data movement. (a) Fixing the outer index

o

selects the slice

Z_{o,:,:}

, in which the query pair

(o,q)

attends to the pairs

(o,k)

. (b) A preprocessing kernel computes the compensated row correction

\hat{\delta}

in FP32 and writes it to HBM. (c) Each thread block of the main backward kernel owns one key tile and one outer group, and further thread blocks run concurrently on other SMs. It stages

K

V

, and

M

once and double-buffers

Q

dO

, and the bias tile in shared memory (SMEM), recomputes

S

P

dP

, and

dS

once per query tile, and accumulates

dK

and

dV

in FP32 registers (RMEM). The score gradient passes through SMEM to form one query-gradient partial per key tile and, through a fixed FP32 tree, one pair-bias partial per outer group. (d) The thread block writes

dK

dV

, and its partials to HBM; separate reduction kernels add the partials in ascending index order in FP32, and the pair-bias partial buffers are reused once consumed (dashed arrow).

$N$	$W_{Q}$	$W_{\mathrm{bias,all}}$	$W_{\mathrm{bias,res}}$	$W_{E}$	$W_{dS}$
384	216	108	108	36	864
640	1000	500	500	100	4000
768	1728	864	252	144	6912
1024	4096	2048	256	256	16384

Appendix

Table 2: Analytical allocation sizes in MiB for

B=1

N_{o}=N

H=4

D=32

\tau_{k}=64

G=4

, and BF16 partials. These are tensor sizes, not measured training peaks; crop 1024 is extrapolated.

Backend	Auto-	$P$	No dense	No $dQ$	No bias	Fixed
	tune	recomputations	FP32 $dS$	atomics	atomics	order
TurboPairFormer	No	1	Yes	Yes	Yes	Yes
TriFast	Yes	3	Yes	Yes	Yes	Yes
Protenix	Yes	3	Yes	Yes	Yes	Yes
MegaFold	No	2	Yes	Yes	No	No
OpenFold3 Triton (newer)	No	2	Yes	No	Yes	No

Appendix

Table 3: Structure of the inspected backward kernels. The three “No” columns indicate the absence of dense FP32 score-gradient storage or floating-point atomic accumulation. “Fixed” denotes an accumulation order independent of block retirement under fixed execution settings. Partial storage in TurboPairFormer remains cubic in

N

Backend	Crop 384	Crop 640	Crop 768
TurboPairFormer	7.01	11.17	14.67
Protenix	7.35	11.30	14.92
cuEquivariance	7.43	12.11	16.51
TriFast	7.51	11.44	15.14
MegaFold	7.63	13.11	18.02
DeepSpeed-Evo	7.52	13.34	18.42

Appendix

Table 4: Mean seconds per optimizer step on eight NVIDIA H200 GPUs, excluding 20 warmup steps from each 400-step run. Protenix, TriFast, MegaFold, and ScaleFold use cuEquivariance triangle multiplication. Bold denotes the lowest mean in each column.

Figure 8: Training stability after a shared 1,000-step learning-rate warmup. Five backend configurations resume the same checkpoint on one H100, with the same restart seed and matched sample order within each crop. Curves show trailing 25-update mean losses, except OpenFold3 Triton, whose finite raw losses are shown before the first NaN. MegaFold curves at crops 384 and 768 exceed the displayed range; 53.9 is the maximum raw loss at crop 768.

Setting	Post-warmup study	Fine-tuning validation
Starting weights	Preview2 after 1,000 updates at crop 640	Preview2
GPUs	One H100	16 H100 (two nodes)
Crop size	384, 640, 768	384
Training updates	400	1,000
Learning rate	$1.8\times 10^{-3}$ , post-warmup plateau	$3\times 10^{-4}$ , 50 warmup steps
EMA decay	0.999	0.98

Appendix

Table 5: Settings of the post-warmup study (Figure 8 ) and of fine-tuning (Section 3.4 ).

Figure 9: Residual compensation on and off. TurboPairFormer runs with and without the saved output residual on one H100, with identical inputs and incoming gradients. Bars show the mean absolute RMSE of the public BF16 gradients against an FP64 reference over ten seeds, and error bars the sample standard deviation. Dashed lines mark the RMSE of the FP64 reference rounded once to BF16. The residual enters only the softmax correction, so

dV

is bitwise identical in both settings.

Quantity	Protenix	cuEquivariance	TurboPairFormer
$O$	0.1695	0.1686	0.1686
$dQ$	0.8202	0.8222	0.3769
$dK$	0.4347	0.4358	0.2437
$dV$	0.2201	0.2164	0.2164
$d\,\mathrm{bias}$	0.3534	0.3506	0.2449

Appendix

Table 6: Median relative

L_{2}

error (%) against an FP64 reference over 32 captured OpenFold3 attention calls. Output error is evaluated on valid query rows. Bold marks the lowest value in each row.

Figure 10: Bitwise gradient replay across eight attention implementations. Each shape contains 100 fixed synthetic inputs, with five calls per input on one H100. A case passes only when all four public BF16 gradients are finite and bitwise identical across all five calls; forward output is excluded from this joint criterion.

Checkpoint	Difference	95% interval
Preview2 (step 0)	$-0.0076$	$[-0.0236,+0.0052]$
PyTorch	$-0.0030$	$[-0.0092,+0.0022]$
OpenFold3 Triton	$+0.0012$	$[-0.0036,+0.0057]$
cuEquivariance	$-0.0035$	$[-0.0092,+0.0024]$
Protenix	$-0.0009$	$[-0.0083,+0.0061]$
TriFast	$-0.0030$	$[-0.0094,+0.0032]$

Appendix

Table 7: Paired model-selection difference, TurboPairFormer minus each checkpoint, with targets averaged over the five seeds and 95% bootstrap intervals over the 85 targets.

$D$	$N$	Ours	cuEq	Protenix	Ours	cuEq	Protenix	Ours	cuEq	Protenix	Fwd	Bwd	F+B
		Forward			Backward			Forward + backward			Change vs. cuEq (%)
16	384	0.384	0.640	0.417	1.167	2.097	1.285	1.493	2.623	1.587	$-39.9$	$-44.3$	$-43.1$
16	640	1.177	2.041	1.316	4.300	8.046	4.655	5.398	9.918	5.830	$-42.3$	$-46.6$	$-45.6$
16	768	1.870	3.341	2.137	7.392	13.244	7.563	9.025	16.478	9.499	$-44.0$	$-44.2$	$-45.2$
32	384	0.474	0.708	0.502	1.520	2.430	1.482	1.920	3.015	1.854	$-33.1$	$-37.5$	$-36.3$
32	640	1.538	2.192	1.602	5.468	8.764	5.654	6.810	10.823	7.129	$-29.8$	$-37.6$	$-37.1$

Appendix

Table 8: Time in ms of one attention call on one H100 (

B=1

N_{o}=N

H=4

) for TurboPairFormer (Ours) in its training configuration, cuEquivariance 0.10.0 (cuEq), and Protenix after autotuning. The last three columns give TurboPairFormer’s change in time relative to cuEquivariance; negative values mean TurboPairFormer is faster.

	Forward time (ms)			Forward reduction (%)		F+B reduction (%)
$N$	Baseline	Core	Combined	Core	Combined	Core	Combined
384	0.547	0.408	0.388	25.5	29.1	8.1	9.2
640	2.037	1.478	1.426	27.4	30.0	8.3	9.0
768	3.336	2.419	2.341	27.5	29.8	7.8	8.5

Appendix

Table 9: Forward dataflow ablation on one H100 (

B=1

N_{o}=N

H=4

D=32

). Times are means of 160 forward calls per arm; percentages are mean reductions over eight paired blocks relative to the baseline.

		Incoming		Outgoing
$N$	Channels	cuEquivariance	Ours	cuEquivariance	Ours
384	128	2.769	1.326 ( $-52.1\%$ )	2.777	1.323 ( $-52.4\%$ )
640	128	3.615	3.334 ( $-7.8\%$ )	3.678	3.336 ( $-9.3\%$ )
768	128	5.142	4.788 ( $-6.9\%$ )	5.210	4.782 ( $-8.2\%$ )
384	64	2.734	1.111 ( $-59.4\%$ )	2.735	1.116 ( $-59.2\%$ )
640	64	2.777	1.921 ( $-30.8\%$ )	2.784	1.922 ( $-31.0\%$ )

Appendix

Table 10: Median time in ms of one triangle multiplicative update, forward plus backward, on one H100 with one example. Parentheses give TurboPairFormer’s (Ours) change in time relative to cuEquivariance 0.10.0; negative values mean TurboPairFormer is faster.

$G$ outer indices per thread block
Thread block owns	Local	Cross-block	Partials	Thread blocks
Key tile (ours)	$dK$ , $dV$	$dQ$ , $d\,\mathrm{bias}$	$J_{K}$	$HJ_{K}N_{o}/G$
Query tile	$dQ$	$dK$ , $dV$ , $d\,\mathrm{bias}$	$2J_{Q}$	$HJ_{Q}N_{o}/G$
Whole slices (DeepSpeed-Evo)	$dQ$ , $dK$ , $dV$	$d\,\mathrm{bias}$	none	$HN_{o}/G$
All $N_{o}$ outer indices per thread block
Query–key tile	$d\,\mathrm{bias}$	$dQ$ , $dK$ , $dV$	$J_{K}+2J_{Q}$	$HJ_{Q}J_{K}$

Appendix

Table 11: Work partitioning for one example. A thread block handles one head, owns the listed part of the

(o,q,k)

tile space, and iterates over the axes it spans.

J_{Q}

and

J_{K}

are the numbers of query and key tiles. The partials column counts the

Q

-sized partial tensors needed when the cross-block reductions of slice-shaped gradients use materialized partials; the upper block also needs

L=N_{o}/G

bias-sized partials for the pair bias.

Figure 11: Making the cross-block reductions deterministic. (a) Four thread blocks (CTA 0–3) contribute to one query-gradient reduction; the bottom row gives each scheme’s consequence for the pair-bias gradient, which sums across slices. With semaphore-ordered atomic adds, the blocks add into one FP32 buffer in turns and wait for them; FlashAttention-3 grants the turns in key-tile order, and DASH chooses the order together with the schedule. With one thread block per slice, a single block visits all

J_{K}

key tiles

t

of the slice, so the launch has

J_{K}

times fewer thread blocks; the deterministic mode of FlashAttention-2 falls in this class here. With materialized partials, blocks write one partial per key tile without waiting and a reduction kernel sums them in index order; TurboPairFormer uses this for

dQ

(

J_{K}

partials) and

d\,\mathrm{bias}

(

L

partials). (b) Change in backward time of variants of our kernel relative to the deterministic kernel at

D=32

on one H100 (Table 12 and Section I.4 ). The atomic variants are not deterministic; “both atomic” combines TMA bulk reduction for

dQ

with atomic

d\,\mathrm{bias}

. Semaphore-ordered variants are not measured.

			Change with non-deterministic atomic accumulation (%)
$D$	$N$	Deterministic (ms)	Atomic $d\,\mathrm{bias}$	Per-element $dQ$	Both, per-element	TMA bulk $dQ$	Both, TMA bulk
16	384	0.849	$-8.0$	$+1.9$	$-6.7$	$-5.2$	$-11.0$
16	640	3.582	$-7.3$	$-0.3$	$-7.2$	$-7.1$	$-12.6$
16	768	6.349	$-7.9$	$-0.1$	$-7.8$	$-7.2$	$-13.7$
32	384	1.091	$-4.1$	$+21.5$	$+13.6$	$+8.3$	$+2.2$
32	640	4.505	$-2.8$	$+18.9$	$+13.3$	$+4.6$	$-0.1$

Appendix

Table 12: Backward time of one TurboPairFormer attention call on one H100 (

B=1

N_{o}=N

H=4

). The deterministic kernel reduces both gradients through materialized partials. Each other column replaces the named reductions with non-deterministic atomic accumulation in FP32 and gives the change in time relative to the deterministic kernel; positive values mean the variant is slower. “Both” combines atomic pair-bias accumulation with the named query-gradient accumulation.

Figure 12: Outgoing triangle multiplicative update. The fused input kernel normalizes the pair features and computes five independent projections: two branch values, two branch gates, and the output gate logits. The normalized features use shared memory, while branch values remain in registers until gating and masking. A separate cuBLAS batched GEMM contracts the branch operands independently for each channel. The fused output kernel combines normalization, output projection, and multiplication by the sigmoid of the saved output gate logits.

Crop	Attention only	TriMul only	Both
384	2.21	2.86	5.55
640	5.59	2.78	8.45
768	6.43	2.99	9.63

Appendix

Table 13: Training step latency reduction (%) from replacing attention, TriMul, or both in a cuEquivariance-enabled OpenFold3 baseline.

Explore similar work

Apr 28, 2026cs.LG

Block-Wise Differentiable Sinkhorn Attention: Tail-Refinement Gradients with a Gap-Aware Dustbin Bridge

We study long-context balanced entropic optimal transport (OT) attention on TPU hardware through a stopped-base, fixed-depth tail-refinement surrogate. After a stopped

T

-step Sinkhorn solve, we unroll a short refinement tail and differentiate that surrogate exactly. For the reported

R=2

TPU path, the backward pass contains four staircase plan factors. We prove an exact one-reference-tile schedule: the

R=2

score cotangent is a single reference plan tile times an explicit modifier field built from vector cotangents and dual differences. This yields block-wise cost

O((T+R)LW)

O(Ld)

input storage, and

O(L)

additional HBM usage for fixed head dimension

d

and band width

W

on the balanced fixed-support path. We also formalize the current \texttt{dustbin_block} path as the same unit-target surrogate on an augmented support, so the adjoint schedule lifts to the single-active-dustbin path used in our TPU runs; this bridge is algebraic and does not claim a general KL-unbalanced or arbitrary-capacity gap model. We provide a local surrogate-bias bound, an a posteriori bias certificate, and a projective contraction certificate for strictly positive active blocks. On synthetic masked problems, the optimized kernel matches exact autodiff of the same centered surrogate to within

10^{-5}

10^{-10}

. On TPU v6e-8, a four-configuration Pfam screen completes end-to-end, and a promoted balanced

R=2

run sustains roughly

8.5

examples per second through a three-hour budget, reaching step

1437

. Held-out Pfam test shards improve reconstruction from

5.57

2.05

and sparse CE from

5.53

5.30

relative to step

0

, with CE logged diagnostically rather than optimized directly; target-barycenter alignment metrics do not materially improve, and a deterministic diagonal reference remains stronger on those metrics.

Dylan Forde

Independent Researcher

Fused Sinkhorn-Localized Similarity Bridge Surrogates

Ask PDF

Sep 28, 2026cs.LG

Broken Symmetry in BF16 Attention: Why FlashAttention Gradients Blow Up Late in Training

BF16 is now standard in large-scale pretraining, including in fused attention kernels such as FlashAttention, and these kernels are widely trusted. When we used FlashAttention-3 to pretrain a 450M-parameter transformer on 50B tokens, however, we ran into a problem: training was healthy for 25B tokens, then the gradient norm grew a thousandfold and the loss ended 0.2 nats above FP32 attention, without a single NaN. Recomputing the attention backward of just two layers in FP32 removes almost all of the excess gradient. Part of the cause is known: a fused multiply-add in the forward softmax, so far treated as an extreme-input NaN case and never fixed in FlashAttention-3. Repairing it stops the blow-up, but the query gradient is still wrong by more than its own size, and training still drives attention logits to thousands of times their size under accurate gradients. The remaining error comes from a broken conservation law. The softmax score gradient sums to zero along every row, which makes the query gradient blind to where the keys sit as a group; rounding it to BF16 leaves a small nonzero sum that leaks the mean key into the gradient, and the leak grows exactly as late training makes keys large and attention sharp. We introduce GProj (gauge projection), which restores the zero sum after the cast with two rank-one corrections per row. It cuts the remaining median query/key gradient errors from 219%/13% to 0.34%/0.37%, on par with FP32 attention, for 4.7% more time per training step. In matched from-scratch runs it trains to the same loss as FP32 attention, while FlashAttention-3 and key smoothing both destabilize.

Junlin Chen, Daize Dong, Huanwei Di +9

Rutgers University · Carnegie Mellon University · Oracle +2

Block Sparse Flash Attention Matched Fp16 Intermediate Pretraining Symmetry

Ask PDF

Sep 21, 2026cs.LG

FlashBoB: I/O-Efficient Exact Backward-over-Backward for Softmax Attention

Transformer models built on the attention mechanism have become a central building block in modern deep learning, yet softmax attention remains a major bottleneck for long-context workloads. While FlashAttention makes the forward and first backward passes I/O-efficient, it does not support backward-over-backward (BoB), which enables exact differentiation through the backward pass for applications such as second-order optimization, test-time training, gradient-based memory, and meta-learning. Existing BoB implementations either materialize large intermediate tensors or exhaust GPU memory at long sequence lengths. We present FlashBoB, an exact, I/O-efficient algorithm for BoB in softmax attention that keeps computation within on-chip tiles and avoids all

N \times N

intermediate tensors, where

N

is the sequence length. The key insight is a hierarchical affine structure in the softmax double backward: two row-wise scalars determine all outputs through affine transformations. This yields a two-pass schedule with bounded on-chip static random-access memory (SRAM) usage and minimal off-chip high-bandwidth memory (HBM) traffic. FlashBoB achieves

Θ(N^2 d^2/M)

HBM traffic (

d

is the head dimension and

M

is the memory size) and, within the standard FlashAttention-style score-recomputation model, matches the inherited large-cache lower bound for exact forward attention. Empirically, it scales exact attention BoB to

N=262\text{K}

on a single A100 80GB GPU, where prior PyTorch exact baselines fail by

N=16\text{K}

, and is up to

6.3\times

faster than FlashBack. These results make exact second-order attention practical at long-context sequence lengths where prior implementations cannot run efficiently.

Anthony Givans, Michael Crawshaw, Mingrui Liu

Block Sparse Flash Attention Gumbel-Softmax Relaxation Backward Passes Activation Functions Gradient-Based Optimization

Ask PDF