Supporting other CUDA archs via --gpu-arch flag

2026-02-11 13:10:41 -08:00
parent c6b245cfe8
commit 3b1e49dcb0
5 changed files with 47 additions and 13 deletions
--- a/Dockerfile.wheels
+++ b/Dockerfile.wheels
@@ -90,7 +90,10 @@ RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    fi

 # Setup Env for Runtime
-ENV TORCH_CUDA_ARCH_LIST=12.1a
+ARG TORCH_CUDA_ARCH_LIST="12.1a"
+ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}
+ARG FLASHINFER_CUDA_ARCH_LIST="12.1a"
+ENV FLASHINFER_CUDA_ARCH_LIST=${FLASHINFER_CUDA_ARCH_LIST}
 ENV TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
 ENV TIKTOKEN_ENCODINGS_BASE=$VLLM_BASE_DIR/tiktoken_encodings