diff --git a/docker/Dockerfile b/docker/Dockerfile index c1a5fbe18..0989e5561 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -587,14 +587,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \ rm /tmp/requirements-cuda.txt /tmp/common.txt -# Install FlashInfer pre-compiled kernel cache and binaries -# This is ~1.1GB and only changes when FlashInfer version bumps +# Install FlashInfer JIT cache (requires CUDA-version-specific index URL) # https://docs.flashinfer.ai/installation.html # From versions.json: .flashinfer.version ARG FLASHINFER_VERSION=0.6.6 RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system flashinfer-cubin==${FLASHINFER_VERSION} \ - && uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \ + uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \ --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ && flashinfer show-config diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 44b7c3809..fe566db35 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -10,6 +10,7 @@ torchaudio==2.10.0 torchvision==0.25.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version # FlashInfer should be updated together with the Dockerfile flashinfer-python==0.6.6 +flashinfer-cubin==0.6.6 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to # breaking changes in 1.19.0 nvidia-cudnn-frontend>=1.13.0,<1.19.0