[CI/Docker] Clean up redundant flashinfer cubin downloads (#22491)

This commit is contained in:
Mohammad Miadh Angkad
2026-04-12 12:30:41 -07:00
committed by GitHub
parent d4ad30b94c
commit 701a0e0c25
3 changed files with 2 additions and 69 deletions
+1 -4
View File
@@ -302,9 +302,7 @@ FROM torch_deps AS flashinfer_cache
ARG CUDA_VERSION
ARG INSTALL_FLASHINFER_JIT_CACHE
ARG FLASHINFER_VERSION
ARG BUILD_AND_DOWNLOAD_PARALLEL
# Download flashinfer cubins
# Stage jit-cache artifacts into /flashinfer_jit_output for clean COPY later
RUN --mount=type=cache,target=/root/.cache/pip \
case "$CUDA_VERSION" in \
@@ -319,8 +317,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
python3 -m pip install flashinfer-jit-cache==${FLASHINFER_VERSION} --index-url https://flashinfer.ai/whl/cu${CUINDEX} \
&& cp -r /usr/local/lib/python3.12/dist-packages/flashinfer_jit_cache /flashinfer_jit_output/ \
&& cp -r /usr/local/lib/python3.12/dist-packages/flashinfer_jit_cache-*.dist-info /flashinfer_jit_output/ ; \
fi \
&& FLASHINFER_CUBIN_DOWNLOAD_THREADS=${BUILD_AND_DOWNLOAD_PARALLEL} FLASHINFER_LOGGING_LEVEL=warning python3 -m flashinfer --download-cubin
fi
########################################################
# PARALLEL STAGE 4: Dev Tools Builder (starts from base)