Build and publish cuda enabled docker image #291
Conversation
…parated compute cap (e.g. `CUDA_ARCHS=89,90,120`)
…and allow multi compute caps
han0110
force-pushed
the
han/feature/cuda-enabled-docker-image
branch
from
27715ca to
c5cb7ee
Compare
han0110
force-pushed
the
han/feature/cuda-enabled-docker-image
branch
from
c934a73 to
c4fcd31
Compare
han0110
force-pushed
the
han/feature/cuda-enabled-docker-image
branch
from
c4fcd31 to
39cd4e2
Compare
jsign
left a comment
jsign
left a comment
There was a problem hiding this comment.
LGMT, left some comments for your consideration
|
|
||
| pub use error::Error; | ||
|
|
||
| /// Applies per-zkVM CUDA architecture build args to a Docker build command. |
There was a problem hiding this comment.
Maybe obvious question to understand the intention of the PR: if a machine signals support for 89 and 120, why would they care about 89 and not always select the highest? (as Zisk does apparently)
I was thinking of multi-GPU machines with different kinds of GPUs, but that would mean Zisk would work there.
| .map(|arch| format!("--generate-code arch=compute_{arch},code=sm_{arch} ")) | ||
| .collect::<String>(); | ||
| cmd.build_arg("NVCC_APPEND_FLAGS", flags.trim_end()) |
There was a problem hiding this comment.
I think we could avoid the trailing space with trim_end(), by collecting to String and then join(" "). But NBD.
| .split(',') | ||
| .filter_map(|s| s.parse::<u32>().ok()) | ||
| .max() | ||
| .unwrap_or(120); |
There was a problem hiding this comment.
Should we have this default in some constant and applied to the other zkVMs too? Or mabye there is a reason for this exception in Zisk? (or maybe error, since we might have non-empty expectation from the other method that parses ENV or detect GPU capabilities?)
|
|
||
| # Per-zkVM CUDA architecture translation | ||
| if [ "$CUDA" = true ] && [ -n "$CUDA_ARCHS" ]; then | ||
| case "$ZKVM" in |
| --registry ${{ needs.image_meta.outputs.registry }} \ | ||
| --tag ${{ needs.image_meta.outputs.sha_tag }}-cuda \ | ||
| --base \ | ||
| --cuda-archs '89,120' |
There was a problem hiding this comment.
We have these 89,120 in some places in this files, worth defining a const/env?
2 participants
CUDA_ARCHtoCUDA_ARCHSto accept comma-separated numeric part of CUDA compute capabilites e.g.CUDA_ARCHS=89,120, if the zkVM supports multiple CUDA compute capabilites it'll be forwarded (Airbender, OpenVM, Risc0), if not the largest will be used (ZisK)CUDA_ARCHSto that when building imageCUDA_ARCHS=89,120(L40S, RTX 40 series, RTX 50 series), this works fine for most zkVMs but takes a bit long for Risc0 (~1hr), but it'd be cached in future PR if the dockerfile is untouched.