From f28fa268c9a5d498955fabf19991a1a6bc737772 Mon Sep 17 00:00:00 2001
From: han0110 <tinghan0110@gmail.com>
Date: Tue, 10 Feb 2026 12:29:41 +0000
Subject: [PATCH] feat: update dockerfiles to use the env var that the zkVM
 recognize, and allow multi compute caps

---
 docker/airbender/Dockerfile.base | 7 ++-----
 docker/openvm/Dockerfile.base    | 9 ---------
 docker/openvm/Dockerfile.server  | 3 +++
 docker/risc0/Dockerfile.base     | 6 ++----
 4 files changed, 7 insertions(+), 18 deletions(-)

diff --git a/docker/airbender/Dockerfile.base b/docker/airbender/Dockerfile.base
index dabcf09..00b3a31 100644
--- a/docker/airbender/Dockerfile.base
+++ b/docker/airbender/Dockerfile.base
@@ -8,11 +8,8 @@ RUN rustup default nightly-2025-07-25
 # Whether to enable CUDA feature or not.
 ARG CUDA
 
-# Default to build for RTX 50 series
-ARG CUDA_ARCH=sm_120
-
-# Env read by Airbender crate `gpu_prover`, only taking the numeric part
-ARG CUDAARCHS=${CUDA_ARCH#sm_}
+# Env read by Airbender crate `gpu_prover`, numeric cuda arch IDs (e.g. "120" or "89;120")
+ARG CUDAARCHS=120
 
 # Copy the Airbender SDK installer script from the workspace context
 COPY --chmod=755 scripts/sdk_installers/install_airbender_sdk.sh /tmp/install_airbender_sdk.sh
diff --git a/docker/openvm/Dockerfile.base b/docker/openvm/Dockerfile.base
index b9025de..b4dd56a 100644
--- a/docker/openvm/Dockerfile.base
+++ b/docker/openvm/Dockerfile.base
@@ -8,15 +8,6 @@ FROM $BASE_IMAGE
 # Set default toolchain to nightly
 RUN rustup default nightly
 
-# Whether to enable CUDA feature or not.
-ARG CUDA
-
-# Default to build for RTX 50 series
-ARG CUDA_ARCH=sm_120
-
-# Env variable read by OpenVM crate `cuda-builder`, need to persist it for building `ere-openvm`.
-ENV CUDA_ARCH=${CUDA_ARCH#sm_}
-
 # Copy the OpenVM SDK installer script from the workspace context
 COPY --chmod=755 scripts/sdk_installers/install_openvm_sdk.sh /tmp/install_openvm_sdk.sh
 
diff --git a/docker/openvm/Dockerfile.server b/docker/openvm/Dockerfile.server
index 10e0bba..9412cfa 100644
--- a/docker/openvm/Dockerfile.server
+++ b/docker/openvm/Dockerfile.server
@@ -14,6 +14,9 @@ WORKDIR /ere
 ARG CUDA
 ARG RUSTFLAGS
 
+# Env variable read by OpenVM crate `openvm-cuda-builder`, comma-separated numeric arch IDs (e.g. "120" or "89,120")
+ARG CUDA_ARCH=120
+
 RUN cargo build --release --package ere-server --bin ere-server --features openvm${CUDA:+,cuda} \
     && mkdir bin && mv target/release/ere-server bin/ere-server \
     && cargo clean && rm -rf $CARGO_HOME/registry/
diff --git a/docker/risc0/Dockerfile.base b/docker/risc0/Dockerfile.base
index d3f81c9..ca6595b 100644
--- a/docker/risc0/Dockerfile.base
+++ b/docker/risc0/Dockerfile.base
@@ -13,10 +13,8 @@ RUN [ -n "$CUDA" ] && \
     (curl -o protoc.zip -L https://github.com/protocolbuffers/protobuf/releases/download/v31.1/protoc-31.1-linux-x86_64.zip \
     && unzip protoc.zip -d /usr/local) || true
 
-# Default to build for RTX 50 series
-ARG CUDA_ARCH=sm_120
-ARG NVCC_APPEND_FLAGS="--generate-code arch=compute_${CUDA_ARCH#sm_},code=sm_${CUDA_ARCH#sm_}"
-ENV NVCC_APPEND_FLAGS=${NVCC_APPEND_FLAGS}
+# Fully formed NVCC flags for CUDA arch targeting (e.g. "--generate-code arch=compute_120,code=sm_120")
+ARG NVCC_APPEND_FLAGS="--generate-code arch=compute_120,code=sm_120"
 
 # Copy and run the Risc0 SDK installer script
 COPY --chmod=755 scripts/sdk_installers/install_risc0_sdk.sh /tmp/install_risc0_sdk.sh