feat: support multi cuda archs for zisk

.github/scripts/build-image.sh

@@ -139,10 +139,7 @@ fi
 
 # Default CUDA_ARCHS when --cuda is set but --cuda-archs not specified
 if [ "$CUDA" = true ] && [ -z "$CUDA_ARCHS" ]; then
-    case "$ZKVM" in
-        zisk) CUDA_ARCHS="120" ;;    # Default to RTX 50 series (ZisK only support setting single CUDA arch)
-        *)    CUDA_ARCHS="89,120" ;; # Default to RTX 40 and 50 series
-    esac
+    CUDA_ARCHS="89,120" # Default to RTX 40 and 50 series
 fi
 
 # Per-zkVM CUDA architecture translation
@@ -168,14 +165,9 @@ if [ "$CUDA" = true ] && [ -n "$CUDA_ARCHS" ]; then
             SERVER_ZKVM_BUILD_ARGS+=(--build-arg "NVCC_APPEND_FLAGS=$NVCC_APPEND_FLAGS")
             ;;
         zisk)
-            IFS=',' read -ra ARCH_ARRAY <<< "$CUDA_ARCHS"
-            if [ "${#ARCH_ARRAY[@]}" -ne 1 ]; then
-                echo "Error: Multiple CUDA architectures are not supported for zisk: $CUDA_ARCHS"
-                exit 1
-            fi
-            BASE_ZKVM_BUILD_ARGS+=(--build-arg "CUDA_ARCH=sm_${ARCH_ARRAY[0]}")
-            SERVER_ZKVM_BUILD_ARGS+=(--build-arg "CUDA_ARCH=sm_${ARCH_ARRAY[0]}")
-            CLUSTER_ZKVM_BUILD_ARGS+=(--build-arg "CUDA_ARCH=sm_${ARCH_ARRAY[0]}")
+            BASE_ZKVM_BUILD_ARGS+=(--build-arg "CUDA_ARCHS=$CUDA_ARCHS")
+            SERVER_ZKVM_BUILD_ARGS+=(--build-arg "CUDA_ARCHS=$CUDA_ARCHS")
+            CLUSTER_ZKVM_BUILD_ARGS+=(--build-arg "CUDA_ARCHS=$CUDA_ARCHS")
             ;;
         *)
             ;;

.github/workflows/build-and-push-images.yml

@@ -95,9 +95,6 @@ jobs:
           - risc0
           - sp1
           - zisk
-        include:
-          - zkvm: zisk
-            cuda_archs: '120'
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -119,7 +116,7 @@ jobs:
             --registry ${{ needs.image_meta.outputs.registry }} \
             --tag ${{ needs.image_meta.outputs.sha_tag }}-cuda \
             --base \
-            --cuda-archs '${{ matrix.cuda_archs || env.CUDA_ARCHS }}'
+            --cuda-archs '${{ env.CUDA_ARCHS }}'
 
       - name: Push ere-base and ere-base-${{ matrix.zkvm }} images with CUDA enabled
         run: |
@@ -132,7 +129,7 @@ jobs:
             --registry ${{ needs.image_meta.outputs.registry }} \
             --tag ${{ needs.image_meta.outputs.sha_tag }}-cuda \
             --server \
-            --cuda-archs '${{ matrix.cuda_archs || env.CUDA_ARCHS }}'
+            --cuda-archs '${{ env.CUDA_ARCHS }}'
 
       - name: Push ere-server-${{ matrix.zkvm }} image with CUDA enabled
         run: |
@@ -150,9 +147,6 @@ jobs:
       matrix:
         zkvm:
           - zisk
-        include:
-          - zkvm: zisk
-            cuda_archs: '120'
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -174,7 +168,7 @@ jobs:
             --registry ${{ needs.image_meta.outputs.registry }} \
             --tag ${{ needs.image_meta.outputs.sha_tag }}-cuda \
             --cluster \
-            --cuda-archs '${{ matrix.cuda_archs || env.CUDA_ARCHS }}'
+            --cuda-archs '${{ env.CUDA_ARCHS }}'
 
       - name: Push ere-cluster-${{ matrix.zkvm }} image with CUDA enabled
         run: |

.github/workflows/test-zkvm-zisk.yml

@@ -17,7 +17,6 @@ jobs:
     with:
       zkvm: zisk
       cuda: true
-      cuda_archs: '120'
       cluster: true
       test_threads: 1
       skip_prove_test: true

.github/workflows/test-zkvm.yml

@@ -12,11 +12,6 @@ on:
         required: false
         type: boolean
         default: false
-      cuda_archs:
-        description: 'Comma-separated CUDA archs to gencode'
-        required: false
-        type: string
-        default: '89,120'
       cluster:
         description: 'Whether to build cluster image'
         required: false
@@ -39,6 +34,7 @@ on:
 
 env:
   CARGO_TERM_COLOR: always
+  CUDA_ARCHS: '89,120'
 
 jobs:
   image_meta:
@@ -111,7 +107,7 @@ jobs:
             --registry ${{ needs.image_meta.outputs.image_registry }} \
             --tag ${{ needs.image_meta.outputs.image_tag }}-cuda \
             --cached-tag "$CACHED_TAG" \
-            --cuda-archs '${{ inputs.cuda_archs }}'
+            --cuda-archs '${{ env.CUDA_ARCHS }}'
 
       - name: Build ere-server-${{ inputs.zkvm }} image with CUDA enabled
         run: |
@@ -120,7 +116,7 @@ jobs:
             --registry ${{ needs.image_meta.outputs.image_registry }} \
             --tag ${{ needs.image_meta.outputs.image_tag }}-cuda \
             --server \
-            --cuda-archs '${{ inputs.cuda_archs }}'
+            --cuda-archs '${{ env.CUDA_ARCHS }}'
 
       - name: Build ere-cluster-${{ inputs.zkvm }} image with CUDA enabled
         if: ${{ inputs.cluster && needs.image_meta.outputs.dockerfile_changed == 'true' }}
@@ -130,7 +126,7 @@ jobs:
             --registry ${{ needs.image_meta.outputs.image_registry }} \
             --tag ${{ needs.image_meta.outputs.image_tag }}-cuda \
             --cluster \
-            --cuda-archs '${{ inputs.cuda_archs }}'
+            --cuda-archs '${{ env.CUDA_ARCHS }}'
 
   clippy_via_docker:
     name: Clippy via Docker

crates/dockerized/src/prover.rs

@@ -37,7 +37,7 @@ pub use error::Error;
 /// - Airbender: `CUDAARCHS` (semicolon-separated, e.g. "89;120")
 /// - OpenVM: `CUDA_ARCH` (comma-separated, e.g. "89,120")
 /// - Risc0: `NVCC_APPEND_FLAGS` (nvcc --generate-code flags)
-/// - Zisk: `CUDA_ARCH` (support only one CUDA architecture, e.g. "sm_120")
+/// - Zisk: `CUDA_ARCHS` (comma-separated, e.g. "89,120")
 fn apply_cuda_build_args(
     cmd: DockerBuildCmd,
     zkvm_kind: zkVMKind,
@@ -74,14 +74,12 @@ fn apply_cuda_build_args(
             cmd.build_arg("NVCC_APPEND_FLAGS", value)
         }
         zkVMKind::Zisk => {
-            if cuda_archs.len() != 1 {
-                return Err(Error::UnsupportedMultiCudaArchs(
-                    zkVMKind::Zisk,
-                    cuda_archs.to_vec(),
-                ));
-            }
-            let value = format!("sm_{}", cuda_archs[0]);
-            cmd.build_arg("CUDA_ARCH", value)
+            let value = cuda_archs
+                .iter()
+                .map(|arch| arch.to_string())
+                .collect::<Vec<_>>()
+                .join(",");
+            cmd.build_arg("CUDA_ARCHS", value)
         }
         _ => cmd,
     })

crates/dockerized/src/prover/error.rs

@@ -1,6 +1,5 @@
 use core::time::Duration;
 
-use ere_catalog::zkVMKind;
 use ere_prover_core::CommonError;
 use ere_server_client::{TwirpErrorResponse, url};
 use thiserror::Error;
@@ -24,10 +23,6 @@ pub enum Error {
     CommonError(#[from] CommonError),
     #[error(transparent)]
     ParseUrl(#[from] url::ParseError),
-    #[error(
-        "Multiple CUDA architectures are not supported for {0:?}, CUDA_ARCHS set or detected: {1:?}"
-    )]
-    UnsupportedMultiCudaArchs(zkVMKind, Vec<u32>),
     #[error("zkVM method error: {0}")]
     zkVM(String),
     #[error("Connection to zkVM server timeout after 5 minutes")]

docker/zisk/Dockerfile.cluster

@@ -54,8 +54,8 @@ RUN git clone https://github.com/han0110/zisk.git --depth 1 --branch patch/v0.16
 # 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 ZisK crate `proofman-starks-lib-c`, comma-separated numeric arch IDs (e.g. "120" or "89,120")
+ARG CUDA_ARCHS=120
 
 # Build binaries
 RUN cargo build --release ${CUDA:+--features gpu}

docker/zisk/Dockerfile.server

@@ -14,8 +14,8 @@ WORKDIR /ere
 ARG CUDA
 ARG RUSTFLAGS
 
-# Default to build for RTX 50 series
-ARG CUDA_ARCH=sm_120
+# Env variable read by ZisK crate `proofman-starks-lib-c`, comma-separated numeric arch IDs (e.g. "120" or "89,120")
+ARG CUDA_ARCHS=120
 
 RUN cargo build --release --package ere-server --bin ere-server --features zisk${CUDA:+,cuda} \
     && mkdir bin && mv target/release/ere-server bin/ere-server \