NVIDIA presentó un conjunto de herramientas que aceleran de extremo a extremo las tuberías de predicción de estructuras biomoleculares, integradas en el BioNeMo Agent Toolkit. La propuesta cubre tres cuellos de botella clásicos del co-folding (plegamiento conjunto de proteínas): la generación del alineamiento múltiple de secuencias (MSA), la inferencia del modelo y el escalado a varias GPU.

Los números que la compañía pone sobre la mesa son concretos:

  • MMseqs2-GPU logra generar el MSA hasta 177 veces más rápido que JackHMMER en CPU, sobre arquitecturas NVIDIA Hopper y Blackwell.
  • cuEquivariance recorta hasta 3 veces el tiempo de la pasada hacia adelante de OpenFold3 y extiende el límite de secuencia a unos 5.900 tokens.
  • El OpenFold3 NIM alcanza hasta 6.400 tokens en una sola GPU B300.
  • Fold-CP habilita inferencia con paralelismo de contexto y permite modelar ensamblajes de hasta 32.000 tokens en 64 GPU B300, cerca de un salto de 12 veces frente al límite de una sola GPU.

¿Por qué importa acelerar el co-folding?

La predicción de estructuras y el co-folding con modelos como OpenFold3 ya son cargas de trabajo masivas y de uso corriente que impulsan el descubrimiento de fármacos y el diseño de proteínas. Cada vez más, agentes de IA las ejecutan de punta a punta, y para que ese agente corra bien la tubería, cada paso debe ser rápido y escalable. Un cuello de botella en cualquier etapa limita el rendimiento total.

La velocidad y la eficiencia de memoria son críticas en flujos como el virtual screening, donde se evalúan millones o miles de millones de compuestos contra uno o pocos blancos proteicos. Los modelos de co-folding suelen dar las mejores estructuras predichas, pero resultan caros de ejecutar, lo que hasta ahora los volvía poco prácticos para tamizados a gran escala.

El otro gran desafío es el tamaño de los ensamblajes. El tiempo de ejecución de estos modelos crece de forma cúbica respecto al número de residuos, y la memoria de una sola GPU impone un techo duro al tamaño del complejo que se puede predecir de una vez.

Figura 1. El flujo de predicción de estructuras acelerado por NVIDIA: MSA en GPU, cuEquivariance, inferencia optimizada y Fold-CP, orquestados por el BioNeMo Agent Toolkit
Figura 1. El flujo de predicción de estructuras acelerado por NVIDIA: MSA en GPU, cuEquivariance, inferencia optimizada y Fold-CP, orquestados por el BioNeMo Agent Toolkit

Eliminar el cuello de botella del MSA

Para los modelos de co-folding, construir el MSA fue tradicionalmente un paso limitado por la CPU que podía dominar el tiempo total. MMseqs2-GPU traslada la búsqueda de homología a las GPU de NVIDIA y escala con la longitud de la secuencia en arquitecturas Hopper y Blackwell.

La última versión acelerada por GPU suma optimizaciones específicas para Hopper y Blackwell, incluido el soporte eficiente para bases de datos mayores que la memoria de la GPU en sistemas Grace, y aceleraciones adicionales gracias a las instrucciones DPX mejoradas de Blackwell disponibles desde CUDA 13.2. Estas contribuciones se integraron de vuelta al repositorio principal de MMseqs2, de modo que toda la comunidad se beneficie.

El MSA Search NIM usa MMseqs2-GPU, cuyo artículo en Nature Methods reporta hasta 177 veces más velocidad de alineamiento que JackHMMER en CPU sobre una sola L40S. En los benchmarks de NVIDIA, la etapa escala sin sobresaltos más allá de los 10.000 tokens en GPU H100 y B300.

Figura 2. La latencia de la etapa de MSA escala de forma casi lineal hasta más de 10.000 tokens en H100 y B300 (429 s frente a 463 s cerca de los 10.000 tokens)
Figura 2. La latencia de la etapa de MSA escala de forma casi lineal hasta más de 10.000 tokens en H100 y B300 (429 s frente a 463 s cerca de los 10.000 tokens)

Plegar a velocidad de estado del arte

cuEquivariance es una biblioteca CUDA-X de primitivas de aprendizaje geométrico para el modelado atomístico. Ofrece versiones aceleradas de los kernels de Triangle Attention, Triangle Multiplication y Attention Pair Bias que dominan el co-folding. En B300 recorta la latencia hasta 3 veces.

Sus kernels se integran directamente en modelos de código abierto como OpenFold3 (como dependencia opcional), OpenFold2, RosettaFold3, Protenix y Boltz. Como las aceleraciones están incorporadas río arriba, el investigador obtiene la mejora de forma automática con solo correr en una GPU NVIDIA el modelo que ya usa. Además, cuEquivariance extiende la longitud máxima de secuencia a unos 5.900 tokens, mientras que PyTorch se queda sin memoria más allá de los 1.500 a 2.500 tokens.

Figura 3. cuEquivariance baja la latencia y extiende la longitud máxima de secuencia frente a PyTorch en H100 y B300
Figura 3. cuEquivariance baja la latencia y extiende la longitud máxima de secuencia frente a PyTorch en H100 y B300

Sobre cuEquivariance, el OpenFold3 NIM aplica optimizaciones de inferencia adicionales que suman al resultado y alcanzan longitudes de hasta 6.400 tokens en una sola B300.

Figura 4. El OpenFold3 NIM totalmente optimizado en H100 y B300 logra plegar más de 3 y 4 veces más rápido, hasta 6.400 tokens en una sola B300
Figura 4. El OpenFold3 NIM totalmente optimizado en H100 y B300 logra plegar más de 3 y 4 veces más rápido, hasta 6.400 tokens en una sola B300

Escalar más allá de una GPU con Fold-CP

La memoria de una sola GPU limitó históricamente los modelos de co-folding a unos pocos miles de residuos. Fold-CP introduce una técnica de paralelización en la que la memoria por dispositivo escala como O(N2/P), donde N es el número de tokens y P la cantidad de GPU. Así alcanza 32.000 tokens en 64 GPU B300 con el modelo Boltz-2, cerca de un salto de 12 veces sobre el límite de una sola GPU.

Para probar Fold-CP, basta apuntar el agente al repositorio de Boltz-CP y pedirle que ejecute inferencia multi-GPU:

Código
# Inferencia con paralelismo de contexto en 4 GPU con Fold-CP (boltz-cp)
git clone https://github.com/NVIDIA-Digital-Bio/boltz-cp && cd boltz-cp
# Instalar dependencias y luego ejecutar
torchrun --nnodes 1 --nproc_per_node 4 \
  src/boltz/distributed/main.py predict /path/to/preprocessed_data \
  --out_dir ./predictions \
  --size_dp 1 --size_cp 4 \
  --recycling_steps 3 --sampling_steps 200 --diffusion_samples 5

¿Qué problemas nuevos habilita esto?

Las mejoras abren clases de problemas de biología estructural antes inalcanzables. En el virtual screening, una inferencia de co-folding más rápida permite aplicar métodos basados en estructura en etapas mucho más tempranas y sobre bibliotecas de compuestos mucho más grandes, lo que mejora la calidad y diversidad de los candidatos que avanzan.

Para los grandes ensamblajes biomoleculares, la combinación de mayor capacidad en una sola B300 y el marco Fold-CP cambia lo que es modelable: complejos a escala del ribosoma, el espliceosoma o grandes ensamblajes de señalización, que eran intratables para estos modelos. Plegar un complejo de unos 10.000 residuos, cerca del tamaño del ribosoma bacteriano, habría sido prohibitivamente caro o directamente imposible en una sola GPU; con B300 y Fold-CP, esas predicciones se vuelven abordables en un nodo multi-GPU.

Para los laboratorios de LatAm que trabajan con presupuestos acotados de cómputo, la clave es que estas aceleraciones llegan río arriba en modelos de código abierto: el investigador que ya corre OpenFold3 o Boltz sobre una GPU NVIDIA hereda la mejora sin reescribir su código, lo que baja la barrera para sumarse al diseño computacional de fármacos.