RTX PRO 6000 /SM_120 incompatibility

[kpetrov]

Hello,

is it possible to run AF2 on Blackwell cards? Would you be so kind as to provide dockerfile/requirements.txt ?

I tried a dozen of possible combinations between jax, cuda, haiku, etc., but none seem to work (XLA fails to initialise cuda, cannot see card, etc.)

thanks!

[k127soo]

I am also experiencing the same issue on a Blackwell-based workstation.

Host environment
GPU: NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition
Driver: 580.xx
CUDA (driver-reported): 13.0

Symptom
MSA generation succeeds
GPUs are visible inside the container
JAX can see both GPUs: devices: [cuda(id=0), cuda(id=1)]
However, even a minimal JAX GPU operation fails before AlphaFold inference starts

Error
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: ptxas exited with non-zero error code 65280, output: ptxas fatal: Program with .target 'sm_90a' cannot be compiled to future architecture

The same issue can be reproduced inside the container with a minimal JAX GPU test: jax.devices() detects both GPUs, but a simple jnp.ones(...) call fails with the same error.

This seems to indicate that the issue is that the JAX/XLA stack inside the AlphaFold 2 container is generating sm_90a code paths that are not compatible with Blackwell.

Has anyone found a working JAX/jaxlib/XLA combination for AlphaFold 2 on Blackwell GPUs?