Does RF3InferenceEngine.run() support parallelism / batching for multiple AtomArray inputs?

[mintuora]

Hi RF3 team, thanks for the great work on RF3!

I have a question regarding performance and parallelism when evaluating multiple structures in a single run.

In my workflow, I often collect multiple generated structures and pass them together to the inference engine, roughly like this:

get atom array from generated outputs

atom_arrays = [out.atom_array for outs in inputs.values() for out in outs]

rf_results = rf3_engine.run(
    inputs=atom_arrays,
    out_dir=None,
)

I noticed from the documentation that RF3 recommends batch processing to amortize startup costs, which makes a lot of sense. However, I’m not fully clear on how RF3InferenceEngine.run() handles multiple inputs internally.

Specifically, I’d like to understand:

Does RF3InferenceEngine.run() internally batch multiple AtomArray inputs into a single forward pass on the GPU?

Or are they processed sequentially in a Python loop?

Is there any built-in parallelism when multiple AtomArray objects are provided?

e.g. GPU batch parallelism, internal DataLoader workers, or other mechanisms?

Are there recommended settings or configuration options (e.g. batch size, inference batch size, diffusion batch size) to maximize throughput when evaluating many structures at once?

If GPU utilization is low even when passing many AtomArray inputs:

Is the recommended approach to increase batch size?

Or is multi-process / multi-GPU execution expected to be the main way to scale evaluation?

My goal is to efficiently evaluate many generated structures in a single RF3 run without repeatedly paying model initialization and CUDA startup costs.

Any guidance on best practices here would be greatly appreciated. Thanks again!

[nscorley]

Hi! Thanks for the question and glad you're using RF3!

1. If you pass multiple AtomArray objects, at the moment, we will multi-process the data loading (CPU), and then loop the model forward pass on GPU. All current structure prediction models work this way due to memory restrictions.

2. If you are running on multiple GPUs, you can pass num_nodes and devices_per_node arguments when you initialize the RF3InferenceEngine. If you do so then the code will automatically distribute the atom arrays across GPUs.

A couple best-practices:

• Initialize RF3 inference engine once (this pays a large cost)
• If you care about speed, you can safely set n_recycles=5 and diffusion_batch_size=1 which should yield significant speedups
• Scale across as many GPUs/nodes as you have available

I would expect around 2-3s/structure/gpu for 200-300 residues proteins on a modern GPU (L40S, A100, etc.) with this setup.

Let me know of other questions!

[mintuora]

Hi, thanks a lot for the detailed explanation — this is very helpful!

I’m currently running RF3 on a single GPU, and I’ve noticed that the GPU memory is not fully utilized during inference. I was wondering whether, in this case, it would be possible to increase diffusion_batch_size, similar to what can be done in RFdiffusion3 (rfd3), in order to better utilize the available memory and potentially improve throughput.

I understand that multiple objects are currently processed with sequential forward passes on the GPU due to memory constraints, so I’m not sure whether increasing diffusion_batch_size on a single card is supported or recommended in RF3, or whether it might introduce any instability or unintended behavior.

Thanks again for your time and for providing such a great tool!

Best regards