software-engineering/spark-vllm-docker
2
0
Fork 0
Code Issues Actions Packages Projects Releases Wiki Activity

Supporting other CUDA archs via --gpu-arch flag

Browse Source
This commit is contained in:
Eugene Rakhmatulin
2026-02-11 13:10:41 -08:00
parent c6b245cfe8
commit 3b1e49dcb0
5 changed files with 47 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
Dockerfile
View File

@@ -55,7 +55,8 @@ ENV CMAKE_CUDA_COMPILER_LAUNCHER=ccache
WORKDIR $VLLM_BASE_DIR
# 2. Set Environment Variables
ENV TORCH_CUDA_ARCH_LIST=12.1a
ARG TORCH_CUDA_ARCH_LIST="12.1a"
ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}
ENV TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
# =========================================================
@@ -90,7 +91,8 @@ FROM base AS builder
# ======= FlashInfer Build ==========
ENV FLASHINFER_CUDA_ARCH_LIST="12.1a"
ARG FLASHINFER_CUDA_ARCH_LIST="12.1a"
ENV FLASHINFER_CUDA_ARCH_LIST=${FLASHINFER_CUDA_ARCH_LIST}
WORKDIR $VLLM_BASE_DIR
ARG FLASHINFER_REF=main
@@ -286,9 +288,11 @@ RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
uv pip install -U transformers --pre; \
fi
# Setup Env for Runtime
ENV TORCH_CUDA_ARCH_LIST=12.1a
ENV FLASHINFER_CUDA_ARCH_LIST="12.1a"
# Setup environment for runtime
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
ENV PATH=$VLLM_BASE_DIR:$PATH