Add support for pre-release FlashInfer packages in Docker builds

This commit is contained in:
Eugene Rakhmatulin
2025-12-20 23:13:26 -08:00
parent 1b3968fe98
commit dfe426e912
4 changed files with 21 additions and 6 deletions

View File

@@ -61,11 +61,13 @@ RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
# Apply in site-packages
RUN patch -p1 -d /usr/local/lib/python3.12/dist-packages < ${VLLM_BASE_DIR}/fastsafetensors.patch
ARG FLASHINFER_PRE=""
# Install flashinfer helper packages
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
uv pip install --system --break-system-packages flashinfer-python -U --no-deps --index-url https://flashinfer.ai/whl && \
uv pip install --system --break-system-packages flashinfer-cubin --index-url https://flashinfer.ai/whl && \
uv pip install --system --break-system-packages flashinfer-jit-cache --index-url https://flashinfer.ai/whl/cu130
uv pip install --system --break-system-packages ${FLASHINFER_PRE} flashinfer-python -U --no-deps --index-url https://flashinfer.ai/whl && \
uv pip install --system --break-system-packages ${FLASHINFER_PRE} flashinfer-cubin --index-url https://flashinfer.ai/whl && \
uv pip install --system --break-system-packages ${FLASHINFER_PRE} flashinfer-jit-cache --index-url https://flashinfer.ai/whl/cu130
# Setup Env for Runtime
ENV TORCH_CUDA_ARCH_LIST=12.1a