AF3 runtime higher than reported for pool of proteins with 5000 tokens

[karchern]

Hi,

we're planning to run AF3 in a large-scale manner to query PPIs on a proteome scale. To that end we're using H100s on a local SCP managed by SLURM.

We're doing a few test folds with pools of proteins (something between 20-35 proteins), all very close to 5000 tokens.

We're seeing very consistent times for this process of around 45 minutes, which is around twice as long as the reported runtime for this token size with an H100 in the performance docs.

Could you give us some hints on what potential reasons could be? I assume I/O is not a major consideration here but would putting the json files onto a scratch disk make sense? Any input would be greatly appreciated. Do you think the number of proteins is an issue?

Below are the SLURM configurations we're using. If you need any other info please lets us know.

#!/bin/bash

#SBATCH -p gpu-el8 
#SBATCH -C gpu=H100 
#SBATCH --cpus-per-task=32
#SBATCH -G 1 
#SBATCH --mem-per-gpu 193028
#SBATCH --time=24:00:00
#SBATCH -o slurm.%N.%j.out
#SBATCH -e slurm.%N.%j.err

module load AlphaFold3/3.0.1.20250908_a8ecdb2-foss-2024a-CUDA-12.6.0

[Augustin-Zidek]

Hello, is this inference-only time, or is this end-to-end (i.e. including MSA and templates search)?

If inference-only, I have a suspicion that this might be either the time to produce the PAE json or to write the PAE json to disk, especially if you are writing to a slow (remote) disk.

To test this hypothesis, could you replace structure_full_data_json with an empty string in this place?

alphafold3/src/alphafold3/model/post_processing.py

Line 73 in 22b9ab8

structure_full_data_json = (

I.e. change that to

structure_full_data_json = ''  # Temporary hack for debugging.

• If the slowness is caused by the conversion to JSON, I can optimize that using C++.
• If the slowness is caused by writing the JSON to disk, try running with the newly added flag --compress_large_output_files.

[karchern]

Hi,

thanks for the suggestions! I was referring to inference-only, sorry for not being clear.

To test your suggestions we (a) moved everything to our scratch disk and (b) disabled the JSON conversion of the PAE information. Unfortunately, the times virtually didn't change, and are still quite far off from the times reported in the performance sheet.

We've tested runtimes for different token sizes on our H100/A100 and we find that while H100 is consistently slower than the suggested optimal runtimes, we find that A100 jobs are faster (except for a ~5k token job, where I guess GPU memory runs out and some swapping occurs).

I'm sharing a table below in case this should ever become relevant for someone else. We will for now try to live with these times :) thanks again

[Augustin-Zidek]

I will run a benchmark on our side to check for regressions, but given our continuous benchmarking, I find it unlikely.

A few more questions:

1. Are you running AF3 with --flash_attention_implementation=triton, or are you setting this differently?
2. Are you seeing this for other models as well (i.e. H100 being slower than A100). If so, it is definitely some misconfiguration of your H100 machines.
3. Are the H100 host machines running with the latest NVIDIA drivers?
4. Are you running with the latest AF3 commit?

[saqzhao]

Hi,

1. We are using the default --flash_attention_implementation (ie not setting it), and are only setting --norun_data_pipeline and --run_inference
2. We don't know if this is also true for other machine learning models. But to be clear, H100 does run faster than A100 and L40 in the graph above (indicated by the H100/A100 right-most column, except for the 1024 tokens run, which we suspect there is some sort of underlying overhead currently unaccounted for), it's just that it's not running as fast as reported in the performance docs.
3. We are using the NVIDIA 590.48.01 driver which supports CUDA 13.1 (reported by nvidia-smi)
4. No, we are using release version 3.0.1 with CUDA 12.6.0 (from this commit a8ecdb2)

We also should probably note that we're not using containerization to run alphafold, but instead use alphafold installed using shared modules on our HPC. We don't think that this could negatively influence performance but could be possible.

We will try again with the latest commit + Singularity (Docker not allowed on our HPC) and get back to you.

[karchern]

I just finished testing this (for A100, 1024 tokens) in a containerized environment using singularity and the newest commit. The runtimes are virtually the same (if not slower) for the containerized version.

I am seeing a lot of warnings regarding autotuning cache misses from triton of the sort below. Could those be connected with the reduced performance?

Below an example of how these warnings look like

ShapeDtype(shape=(1024, 1024, 128), dtype=bfloat16, vmap_axes=()), 'weights': BatchedShapeDtype(shape=(128, 2, 512), dtype=bfloat16, vmap_axes=()), 'activation': <PjitFunction of <function silu at 0x7fff8de40f40>>, 'precisio
n': (Precision.DEFAULT, Precision.DEFAULT), 'return_residuals': False})                                                                                                                                                         
W0210 16:46:23.436067 140737350427392 op.py:517] Autotuning cache miss for PallasTritonGatedLinearUnit(config=None, vjp=GatedLinearUnitVjp(config=None, vjp=None)) on NVIDIA A100-PCIE-40GB with key immutabledict({'x': Batched
ShapeDtype(shape=(1024, 384), dtype=bfloat16, vmap_axes=()), 'weights': BatchedShapeDtype(shape=(384, 2, 1536), dtype=bfloat16, vmap_axes=()), 'activation': <PjitFunction of <function silu at 0x7fff8de40f40>>, 'precision': (
Precision.DEFAULT, Precision.DEFAULT), 'return_residuals': False})                                                                                                                                                             
W0210 16:46:23.459434 140737350427392 op.py:517] Autotuning cache miss for PallasTritonGatedLinearUnit(config=None, vjp=GatedLinearUnitVjp(config=None, vjp=None)) on NVIDIA A100-PCIE-40GB with key immutabledict({'x': Batched
ShapeDtype(shape=(1024, 1024, 128), dtype=bfloat16, vmap_axes=((0, 1),)), 'weights': BatchedShapeDtype(shape=(128, 2, 256), dtype=bfloat16, vmap_axes=(None,)), 'activation': <PjitFunction of <function sigmoid at 0x7fff8de40c
20>>, 'precision': (Precision.DEFAULT, Precision.DEFAULT), 'return_residuals': False})
W0210 16:46:23.496870 140737350427392 op.py:517] Autotuning cache miss for PallasTritonGatedLinearUnit(config=None, vjp=GatedLinearUnitVjp(config=None, vjp=None)) on NVIDIA A100-PCIE-40GB with key immutabledict({'x': Batched
ShapeDtype(shape=(1024, 1024, 128), dtype=bfloat16, vmap_axes=((0, 1),)), 'weights': BatchedShapeDtype(shape=(128, 2, 256), dtype=bfloat16, vmap_axes=(None,)), 'activation': <PjitFunction of <function sigmoid at 0x7fff8de40c
20>>, 'precision': (Precision.DEFAULT, Precision.DEFAULT), 'return_residuals': False})                                                                                                                                         
W0210 16:46:23.536557 140737350427392 op.py:517] Autotuning cache miss for PallasTritonFlashAttention(config=None, vjp=PallasTritonFlashAttentionVjp(config=None, vjp=None, dbias_intermediate_dtype=None), use_base2=False, use
_stable_softmax=<class 'tokamax._src.ops.attention.base.AUTO'>) on NVIDIA A100-PCIE-40GB with key immutabledict({'q': BatchedShapeDtype(shape=(128, 1024, 4, 32), dtype=bfloat16, vmap_axes=((0, 1),)), 'k': BatchedShapeDtype(s
hape=(128, 1024, 4, 32), dtype=bfloat16, vmap_axes=((0, 1),)), 'v': BatchedShapeDtype(shape=(128, 1024, 4, 32), dtype=bfloat16, vmap_axes=((0, 1),)), 'bias': BatchedShapeDtype(shape=(1, 4, 1024, 1024), dtype=bfloat16, vmap_a
xes=((0, 1),)), 'dropout_mask': None, 'dropout_rate': 0.0, 'k_indices': None, 'logits_dtype': dtype('float32'), 'logits_scale': 0.17677669529663687, 'logits_soft_cap': None, 'mask': Mask(bool_mask=BatchedShapeDtype(shape=(12
8, 1, 1, 1024), dtype=bool, vmap_axes=(None,)), q_start=None, q_end=None, k_start=None, k_end=None, is_causal=False), 'normalize_output': True, 'paging_info': None, 'precision': (DotAlgorithmPreset.BF16_BF16_F32, DotAlgorith
mPreset.BF16_BF16_F32), 'q_indices': None, 'return_residuals': False})  ``` 

[Augustin-Zidek]

I finished internal testing and I am seeing times consistent with what is reported in the performance documentation. I am investigating the Triton cache misses.

[Augustin-Zidek]

I found out that we are indeed missing a lot of shapes in the Tokamax autotuning cache. We are fixing this with Tokamax folks on their side, I will then pull in the latest Tokamax version. This should hopefully speed things up.

[karchern]

Thanks a lot for your diligent help, Augustin!

I was wondering: If you experience these cache issues on your end, but still see 2x higher speed wrt to us, can we expect this to fix the issue on our end?

I personally am starting to feel it's likely that we have some misconfiguration or other more hardware-related issues on our H100s (I think that mainly because our A100 are actually performing better than on your end - assuming the A100 performance numbers you report are up to date).

[Augustin-Zidek]

I was wondering: If you experience these cache issues on your end, but still see 2x higher speed wrt to us, can we expect this to fix the issue on our end?

It will make things somewhat better, but unlikely to fix. I am glad I found the issue thanks to this though, this will make AF3 faster for everyone so I am grateful for you reporting it.

I personally am starting to feel it's likely that we have some misconfiguration or other more hardware-related issues on our H100s (I think that mainly because our A100 are actually performing better than on your end - assuming the A100 performance numbers you report are up to date).

This is my feeling as well, worth investigating, running some standard benchmarks and diagnostics on the H100 machines.

[ikalvet]

I just finished testing this (for A100, 1024 tokens) in a containerized environment using singularity and the newest commit. The runtimes are virtually the same (if not slower) for the containerized version.

I am seeing a lot of warnings regarding autotuning cache misses from triton of the sort below. Could those be connected with the reduced performance?

Below an example of how these warnings look like

ShapeDtype(shape=(1024, 1024, 128), dtype=bfloat16, vmap_axes=()), 'weights': BatchedShapeDtype(shape=(128, 2, 512), dtype=bfloat16, vmap_axes=()), 'activation': <PjitFunction of <function silu at 0x7fff8de40f40>>, 'precisio
n': (Precision.DEFAULT, Precision.DEFAULT), 'return_residuals': False})                                                                                                                                                         
W0210 16:46:23.436067 140737350427392 op.py:517] Autotuning cache miss for PallasTritonGatedLinearUnit(config=None, vjp=GatedLinearUnitVjp(config=None, vjp=None)) on NVIDIA A100-PCIE-40GB with key immutabledict({'x': Batched
ShapeDtype(shape=(1024, 384), dtype=bfloat16, vmap_axes=()), 'weights': BatchedShapeDtype(shape=(384, 2, 1536), dtype=bfloat16, vmap_axes=()), 'activation': <PjitFunction of <function silu at 0x7fff8de40f40>>, 'precision': (
Precision.DEFAULT, Precision.DEFAULT), 'return_residuals': False})                                                                                                                                                             
W0210 16:46:23.459434 140737350427392 op.py:517] Autotuning cache miss for PallasTritonGatedLinearUnit(config=None, vjp=GatedLinearUnitVjp(config=None, vjp=None)) on NVIDIA A100-PCIE-40GB with key immutabledict({'x': Batched
ShapeDtype(shape=(1024, 1024, 128), dtype=bfloat16, vmap_axes=((0, 1),)), 'weights': BatchedShapeDtype(shape=(128, 2, 256), dtype=bfloat16, vmap_axes=(None,)), 'activation': <PjitFunction of <function sigmoid at 0x7fff8de40c
20>>, 'precision': (Precision.DEFAULT, Precision.DEFAULT), 'return_residuals': False})
W0210 16:46:23.496870 140737350427392 op.py:517] Autotuning cache miss for PallasTritonGatedLinearUnit(config=None, vjp=GatedLinearUnitVjp(config=None, vjp=None)) on NVIDIA A100-PCIE-40GB with key immutabledict({'x': Batched
ShapeDtype(shape=(1024, 1024, 128), dtype=bfloat16, vmap_axes=((0, 1),)), 'weights': BatchedShapeDtype(shape=(128, 2, 256), dtype=bfloat16, vmap_axes=(None,)), 'activation': <PjitFunction of <function sigmoid at 0x7fff8de40c
20>>, 'precision': (Precision.DEFAULT, Precision.DEFAULT), 'return_residuals': False})                                                                                                                                         
W0210 16:46:23.536557 140737350427392 op.py:517] Autotuning cache miss for PallasTritonFlashAttention(config=None, vjp=PallasTritonFlashAttentionVjp(config=None, vjp=None, dbias_intermediate_dtype=None), use_base2=False, use
_stable_softmax=<class 'tokamax._src.ops.attention.base.AUTO'>) on NVIDIA A100-PCIE-40GB with key immutabledict({'q': BatchedShapeDtype(shape=(128, 1024, 4, 32), dtype=bfloat16, vmap_axes=((0, 1),)), 'k': BatchedShapeDtype(s
hape=(128, 1024, 4, 32), dtype=bfloat16, vmap_axes=((0, 1),)), 'v': BatchedShapeDtype(shape=(128, 1024, 4, 32), dtype=bfloat16, vmap_axes=((0, 1),)), 'bias': BatchedShapeDtype(shape=(1, 4, 1024, 1024), dtype=bfloat16, vmap_a
xes=((0, 1),)), 'dropout_mask': None, 'dropout_rate': 0.0, 'k_indices': None, 'logits_dtype': dtype('float32'), 'logits_scale': 0.17677669529663687, 'logits_soft_cap': None, 'mask': Mask(bool_mask=BatchedShapeDtype(shape=(12
8, 1, 1, 1024), dtype=bool, vmap_axes=(None,)), q_start=None, q_end=None, k_start=None, k_end=None, is_causal=False), 'normalize_output': True, 'paging_info': None, 'precision': (DotAlgorithmPreset.BF16_BF16_F32, DotAlgorith
mPreset.BF16_BF16_F32), 'q_indices': None, 'return_residuals': False})  ``` 

Wanted to report that with upgrading to tokamax 0.0.10, I also see these warnings flooding the STDOUT. Both with H200, A4000 and RTX PRO 4000 GPU's. Just tried the latest tokamax commit, and the issue of flooding cache miss messages still persists. Thank you for digging deep into this though and looking forward to a fix for this!
I don't see any significant performance issues on my ~300-token predictions though. The only noticeable difference (that appeared already when switching to tokamax 0.0.8) was that now the JAX compilation takes slightly longer (90s now, 60s before without tokamax), and somehow the first two predictions seem to spend time compiling (before only the first one took longer). Also those autotuning cache miss warnings are printed only at predictions during which the model is compiled.