From 3af90daf087878c2e8441b3a061c571fd62da2b9 Mon Sep 17 00:00:00 2001 From: han0110 Date: Tue, 21 Apr 2026 10:59:06 +0000 Subject: [PATCH 1/2] feat: use env `CUDA_ARCHS` for zisk to do multi cuda archs codegen --- .github/scripts/build-image.sh | 16 ++++------------ .github/workflows/build-and-push-images.yml | 12 +++--------- .github/workflows/test-zkvm-zisk.yml | 1 - crates/dockerized/src/prover.rs | 16 +++++++--------- crates/dockerized/src/prover/error.rs | 5 ----- docker/zisk/Dockerfile.cluster | 4 ++-- docker/zisk/Dockerfile.server | 4 ++-- 7 files changed, 18 insertions(+), 40 deletions(-) diff --git a/.github/scripts/build-image.sh b/.github/scripts/build-image.sh index 63ef8f6f..8913024a 100755 --- a/.github/scripts/build-image.sh +++ b/.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") ;; *) ;; diff --git a/.github/workflows/build-and-push-images.yml b/.github/workflows/build-and-push-images.yml index ebe005ee..8eeb1092 100644 --- a/.github/workflows/build-and-push-images.yml +++ b/.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: | diff --git a/.github/workflows/test-zkvm-zisk.yml b/.github/workflows/test-zkvm-zisk.yml index 7dea64a3..9ed738df 100644 --- a/.github/workflows/test-zkvm-zisk.yml +++ b/.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 diff --git a/crates/dockerized/src/prover.rs b/crates/dockerized/src/prover.rs index 1e6b4103..4a45d9b1 100644 --- a/crates/dockerized/src/prover.rs +++ b/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::>() + .join(","); + cmd.build_arg("CUDA_ARCHS", value) } _ => cmd, }) diff --git a/crates/dockerized/src/prover/error.rs b/crates/dockerized/src/prover/error.rs index 3fb5e536..261d4734 100644 --- a/crates/dockerized/src/prover/error.rs +++ b/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), #[error("zkVM method error: {0}")] zkVM(String), #[error("Connection to zkVM server timeout after 5 minutes")] diff --git a/docker/zisk/Dockerfile.cluster b/docker/zisk/Dockerfile.cluster index db57f0a2..c5280e29 100644 --- a/docker/zisk/Dockerfile.cluster +++ b/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} diff --git a/docker/zisk/Dockerfile.server b/docker/zisk/Dockerfile.server index aceff040..e19f5174 100644 --- a/docker/zisk/Dockerfile.server +++ b/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 \ From 5a03dbf4a54649e7ba7be0db6968b75bcfcb37c9 Mon Sep 17 00:00:00 2001 From: han0110 Date: Tue, 21 Apr 2026 11:07:16 +0000 Subject: [PATCH 2/2] ci: simplify --- .github/workflows/test-zkvm.yml | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/.github/workflows/test-zkvm.yml b/.github/workflows/test-zkvm.yml index 43c08036..345b02d4 100644 --- a/.github/workflows/test-zkvm.yml +++ b/.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