Profiling ESMFold

Problem

ESMFold predicts a protein's 3D structure from a single amino acid sequence. It replaces the evolutionary database search that AlphaFold relies on with a protein language model. That makes structure prediction cheap enough to run across whole proteomes.

The cost of the alignment is replaced by a cost inside the model. To keep its predictions geometrically consistent, ESMFold maintains a pairwise representation of every residue pair and updates it with triangular attention and triangular multiplication, so memory and compute grow quadratically with sequence length.

I wanted to measure which of the model's stages pays that cost as proteins get longer, and whether standard inference optimizations reduce it.

Profiling

I profiled the forward pass on an H100 across the 33 CASP14 targets (74 to 922 residues) and attributed every kernel to one of five stages (input embedding, bias and MLP, triangular multiplication, triangular attention, structure module). The chart shows how runtime splits between those stages for each target, ordered by sequence length.

The GPU and the CPU were limited by different stages:

• On the GPU, triangular attention takes over as proteins get longer, from about a third of runtime on short sequences to 70% near 1,000 residues. That is the quadratic term appearing exactly where the architecture predicts it.
• On the CPU, time went to triangular multiplication. This stage updates each residue pair from its neighboring pairs through thousands of small elementwise kernels per pass, whereas attention runs as a single fused flash-attention kernel.

Short proteins are therefore bound by kernel launch overhead on the CPU, and long proteins by attention compute on the GPU.

Compiling the model

Graph compilation fuses sequences of small kernels into one, which targets the launch overhead that limits short proteins. The effect varied widely between targets:

• On several short targets it cut the forward pass by about 40%.
• On others the compiled model was slower than the original, by up to 57%.

The traces explained the slowdowns. The folding trunk updates the pairwise representation over several recycling passes and loops over a Python dictionary on each one. The compiler cannot trace that loop, so the graph breaks at that point, control returns to the Python interpreter, and the kernels on either side run unfused. Repeated on every recycling pass, that overhead exceeds what fusion saves, and the compiled model ends up slower than the uncompiled one.

Long proteins barely changed either way, because almost all of their runtime is a single attention kernel that fusion cannot improve. Microbatching the folding trunk in 128-residue chunks did help them. It reduced peak memory enough for a 900-residue protein to fit on a single GPU.

Takeaway

The profile points to three changes, at three different levels:

• Serving: bucket sequences by length, so that short proteins share a batch and amortize their launch overhead while long proteins run alone.
• Kernels: on long proteins the time is inside one attention kernel doing arithmetic, so running the triangular blocks in lower precision (FP8 on H100) is the most direct way to speed them up, provided structure quality holds.
• Model: the pairwise representation gives ESMFold its accuracy, and it is also the reason the targets that matter most in drug discovery (long multi-domain proteins and protein complexes) are the most expensive to predict.

ESMFold has already moved one of AlphaFold's components, the evolutionary search, out of inference and into pretraining. Geometric consistency is the other, and it is still computed at inference time for every residue pair on every pass. In a drug discovery screen that cost recurs for every target, whereas anything learned in training is paid for once.

Geometric consistency could plausibly make the same transition, in one of two ways. Through the loss, by training on distances and angles until the embeddings themselves encode a consistent geometry that a small decoder can turn into coordinates. Or through the vocabulary, by representing local structure as tokens the language model predicts alongside amino acids, so that consistency is learned like grammar: which structural tokens can follow one another. Either way, pairwise computation would be needed only to refine the hardest long-range contacts. The cost of structure prediction would stop growing with the square of protein length, and long targets would become routine at proteome scale.