feat: enable cuequivariance kernels on the base image

Dockerfile

@@ -1,5 +1,5 @@
 # Full performance multi-stage build with complete CUDA toolchain
-FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04 AS builder
+FROM nvidia/cuda:12.6.1-cudnn-devel-ubuntu22.04 AS builder
 
 # Install complete build dependencies including CUDA compiler tools
 RUN apt-get update && apt-get install -y \
@@ -35,7 +35,7 @@ RUN /opt/openfold3/scripts/install_third_party_dependencies.sh
 
 # Install the package
 WORKDIR /opt/openfold3
-RUN python3 setup.py install
+RUN pip install '.[cuequivariance]' --no-build-isolation && python setup.py install
 
 # Set CUDA architecture for compilation (adjust based on your GPU)
 ENV TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"
@@ -45,7 +45,7 @@ ENV TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"
 #     python3 -c "import deepspeed; print('DeepSpeed ops loaded successfully')"
 
 # Runtime stage - use devel image for full CUDA support
-FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04 AS runtime
+FROM nvidia/cuda:12.6.1-cudnn-devel-ubuntu22.04 AS runtime
 
 # Install runtime dependencies
 RUN apt-get update && apt-get install -y \

environments/production.yml

@@ -30,7 +30,9 @@ dependencies:
   - bioconda::kalign2
   - bioconda::snakemake
   - pytorch::pytorch
-  - pytorch::pytorch-cuda
+  - conda-forge::cuda-version=12.6     # ← Force CUDA 12.6
+  - conda-forge::libcublas>=12.6.1     # ← Ensure cuBLAS 12.6.1+
+  # - pytorch::pytorch-cuda
   - memory_profiler
   - func_timeout
   - boto3

openfold3/core/model/layers/triangular_multiplicative_update.py

@@ -30,7 +30,10 @@
 from openfold3.core.model.primitives import LayerNorm, Linear
 from openfold3.core.utils.tensor_utils import permute_final_dims
 
-cueq_is_installed = importlib.util.find_spec("cuequivariance_torch") is not None
+cueq_is_installed = (
+    importlib.util.find_spec("cuequivariance_torch") is not None
+    and torch.cuda.is_available()
+)
 if cueq_is_installed:
     from cuequivariance_torch import triangle_multiplicative_update
 

openfold3/core/model/primitives/attention.py

@@ -46,7 +46,10 @@
 if ds4s_is_installed:
     from deepspeed.ops.deepspeed4science import DS4Sci_EvoformerAttention
 
-cueq_is_installed = importlib.util.find_spec("cuequivariance_torch") is not None
+cueq_is_installed = (
+    importlib.util.find_spec("cuequivariance_torch") is not None
+    and torch.cuda.is_available()
+)
 if cueq_is_installed:
     from cuequivariance_ops_torch.triangle_attention import (
         CUEQ_TRIATTN_FALLBACK_THRESHOLD,

openfold3/tests/compare_utils.py

@@ -30,8 +30,9 @@ def skip_unless_ds4s_installed():
 
 
 def skip_unless_cueq_installed():
-    cueq_is_installed = cueq_is_installed = (
+    cueq_is_installed = (
         importlib.util.find_spec("cuequivariance_torch") is not None
+        and torch.cuda.is_available()
     )
     return unittest.skipUnless(
         cueq_is_installed, "Requires CU-Equivaraince to be installed"