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

@@ -71,11 +71,13 @@ RUN --mount=type=cache,id=pip-cache,target=/root/.cache/pip \
RUN --mount=type=cache,id=pip-cache,target=/root/.cache/pip \
pip install xgrammar fastsafetensors
ARG FLASHINFER_PRE=""
# Install FlashInfer packages
RUN --mount=type=cache,id=pip-cache,target=/root/.cache/pip \
pip install flashinfer-python --no-deps --index-url https://flashinfer.ai/whl --pre && \
pip install flashinfer-cubin --index-url https://flashinfer.ai/whl --pre && \
pip install flashinfer-jit-cache --index-url https://flashinfer.ai/whl/cu130 --pre && \
pip install ${FLASHINFER_PRE} flashinfer-python --no-deps --index-url https://flashinfer.ai/whl && \
pip install ${FLASHINFER_PRE} flashinfer-cubin --index-url https://flashinfer.ai/whl && \
pip install ${FLASHINFER_PRE} flashinfer-jit-cache --index-url https://flashinfer.ai/whl/cu130 && \
pip install apache-tvm-ffi nvidia-cudnn-frontend nvidia-cutlass-dsl nvidia-ml-py tabulate
# =========================================================