[Dep] Upgrade flashinfer to 0.6.14 (#29910)

Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>
Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
This commit is contained in:
Baizhou Zhang
2026-07-09 17:52:33 -07:00
committed by GitHub
co-authored by Brayden Zhong Mohammad Miadh Angkad
parent b2f9a95867
commit 2c6cd1ef41
6 changed files with 48 additions and 25 deletions
+7 -3
View File
@@ -19,7 +19,7 @@ ARG PIP_DEFAULT_INDEX
ARG UBUNTU_MIRROR
ARG GITHUB_ARTIFACTORY=github.com
ARG INSTALL_FLASHINFER_JIT_CACHE=0
ARG FLASHINFER_VERSION=0.6.12
ARG FLASHINFER_VERSION=0.6.14
ARG MOONCAKE_VERSION=0.3.11.post1
ARG MSCCLPP_VERSION=sglang-v0.9.1
#if need other arg please add in MOONCAKE_COMPILE_ARG
@@ -332,7 +332,7 @@ ARG CUDA_VERSION
ARG INSTALL_FLASHINFER_JIT_CACHE
ARG FLASHINFER_VERSION
# Stage jit-cache artifacts into /flashinfer_jit_output for clean COPY later
# Stage jit-cache/cubin artifacts into /flashinfer_jit_output for clean COPY later
RUN --mount=type=cache,target=/root/.cache/pip \
case "$CUDA_VERSION" in \
12.6.1) CUINDEX=126 ;; \
@@ -342,6 +342,10 @@ RUN --mount=type=cache,target=/root/.cache/pip \
*) echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 ;; \
esac \
&& mkdir -p /flashinfer_jit_output \
# flashinfer-cubin is CUDA-version-agnostic, unlike jit-cache, so its index-url has no cu${CUINDEX} suffix
&& python3 -m pip install flashinfer-cubin==${FLASHINFER_VERSION} --index-url https://flashinfer.ai/whl \
&& cp -r /usr/local/lib/python3.12/dist-packages/flashinfer_cubin /flashinfer_jit_output/ \
&& cp -r /usr/local/lib/python3.12/dist-packages/flashinfer_cubin-*.dist-info /flashinfer_jit_output/ \
&& if [ "$INSTALL_FLASHINFER_JIT_CACHE" = "1" ]; then \
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/ \
@@ -458,7 +462,7 @@ COPY --from=deepep_builder /build/DeepEP /sgl-workspace/DeepEP
RUN --mount=type=cache,target=/root/.cache/pip \
pip install /tmp/wheels/deepep/*.whl && rm -rf /tmp/wheels/deepep
# Copy flashinfer jit-cache package (if installed)
# Copy flashinfer cubin (always) and jit-cache (if installed) packages
COPY --from=flashinfer_cache /flashinfer_jit_output/ /usr/local/lib/python3.12/dist-packages/
# Copy dev tools