Segmentation fault / MLIR optimization error on NVIDIA RTX 5090 (Blackwell) with LocalColabFold 1.6.1

[Karol007]

Describe the bug

When trying to run colabfold_batch on a machine equipped with the new NVIDIA RTX 5090 (Blackwell architecture), the process crashes immediately at Query 1 with a Segmentation Fault and MLIR optimization errors. It appears the pre-compiled jaxlib in the current pixi environment does not yet support Blackwell's Compute Capability.

Expected behavior

The model should compile the XLA graph and start predicting structures on the GPU, as it does on Ampere/Ada architectures.

To Reproduce

1. Install localcolabfold using standard pixi instructions.
2. Run standard batch prediction:
   pixi run colabfold_batch msa_batch results --num-models 1 --num-recycle 3
3. Observe crash.

Logs / Error Messages

2026-06-28 18:44:43,611 Running colabfold 1.6.1
2026-06-28 18:44:49,812 Running on GPU
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1782665109.960637 20354 mlir_graph_optimization_pass.cc:425] MLIR V1 optimization pass is not enabled
Segmentation fault (core dumped)

Attempting to disable MLIR (export XLA_FLAGS="--xla_gpu_enable_mlir_graph_optimization=false") results in an unknown flag error.

System details

• OS: Windows 11 WSL2 (Ubuntu 24.04.4 LTS) / Linux
• GPU: NVIDIA RTX 5090 (Blackwell)
• ColabFold version: 1.6.1 (local, via pixi)

Temporary Workaround (for others)

Forcing CPU execution (export JAX_PLATFORMS="cpu" and export CUDA_VISIBLE_DEVICES="") works perfectly, confirming the issue is strictly XLA/GPU related.

Are there any plans to update the jax/jaxlib dependencies to versions compiled with CUDA 12.8+ to support Blackwell?

Best regards
Karol007

[YoshitakaMo]

Hi, @Karol007 .
I just updated the pyproject.toml of LocalColabFold (YoshitakaMo/localcolabfold@e2e8c4a). By introducing jax 0.10.1, I believe I have made it compatible with the NVIDIA RTX 5090. It works for me on my system running Ubuntu 24.04 with an RTX 5090.
Please try installing it using this pyproject.toml. Since this is experimental, I would greatly appreciate your feedback.

I have also added the following to my execution bash script:

#!/bin/bash
...
...
export XLA_FLAGS="--xla_gpu_autotune_level=0"
export TF_ENABLE_ONEDNN_OPTS=0
...
colabfold_batch --amber --templates <input> <output>