chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[chore] Arrange NV packages in Dockerfile (#13749)

Browse Source
This commit is contained in:
Baizhou Zhang
2025-11-27 18:08:27 -05:00
committed by GitHub
parent 4c9f7c97d3
commit 051ad83347
1 changed files with 20 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
docker/Dockerfile
View File

@@ -18,6 +18,7 @@ ARG UBUNTU_MIRROR
ARG GITHUB_ARTIFACTORY=github.com
ARG INSTALL_FLASHINFER_JIT_CACHE=0
ARG FLASHINFER_VERSION=0.5.3
ARG NVSHMEM_VERSION=3.4.5
ENV DEBIAN_FRONTEND=noninteractive \
CUDA_HOME=/usr/local/cuda \
@@ -131,21 +132,16 @@ RUN --mount=type=cache,target=/root/.cache/pip python3 -m pip install --upgrade
&& if [ "$INSTALL_FLASHINFER_JIT_CACHE" = "1" ]; then \
python3 -m pip install flashinfer-jit-cache==${FLASHINFER_VERSION} --index-url https://flashinfer.ai/whl/cu${CUINDEX} ; \
fi \
&& if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
python3 -m pip install nvidia-nccl-cu12==2.28.3 --force-reinstall --no-deps ; \
python3 -m pip install nvidia-cudnn-cu12==9.16.0.29 --force-reinstall --no-deps; \
elif [ "${CUDA_VERSION%%.*}" = "13" ]; then \
python3 -m pip install nvidia-nccl-cu13==2.28.3 --force-reinstall --no-deps ; \
else \
echo "No NCCL mapping for CUDA_VERSION=${CUDA_VERSION}" && exit 1 ; \
fi \
&& FLASHINFER_CUBIN_DOWNLOAD_THREADS=${BUILD_AND_DOWNLOAD_PARALLEL} FLASHINFER_LOGGING_LEVEL=warning python3 -m flashinfer --download-cubin
# Download NVSHMEM source files
# We use Tom's DeepEP fork for GB200 for now; the 1fd57b0276311d035d16176bb0076426166e52f3 commit is https://github.com/fzyzcjy/DeepEP/tree/gb200_blog_part_2
RUN set -eux; \
if [ "${CUDA_VERSION%%.*}" != "13" ]; then \
pip install nvidia-nvshmem-cu12==3.4.5 ; \
if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
pip install nvidia-nvshmem-cu12==${NVSHMEM_VERSION} ; \
elif [ "${CUDA_VERSION%%.*}" = "13" ]; then \
pip install nvidia-nvshmem-cu13==${NVSHMEM_VERSION} ; \
else \
echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 ; \
fi && \
if [ "$GRACE_BLACKWELL" = "1" ]; then \
git clone https://github.com/fzyzcjy/DeepEP.git && \
@@ -198,6 +194,19 @@ RUN --mount=type=cache,target=/root/.cache/pip python3 -m pip install \
nixl \
py-spy
# Some patching packages
# TODO: Remove this when torch version covers these packages
RUN --mount=type=cache,target=/root/.cache/pip if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
python3 -m pip install nvidia-nccl-cu12==2.28.3 --force-reinstall --no-deps ; \
python3 -m pip install nvidia-cudnn-cu12==9.16.0.29 --force-reinstall --no-deps; \
python3 -m pip install nvidia-nvshmem-cu12==${NVSHMEM_VERSION} --force-reinstall --no-deps; \
elif [ "${CUDA_VERSION%%.*}" = "13" ]; then \
python3 -m pip install nvidia-nccl-cu13==2.28.3 --force-reinstall --no-deps ; \
python3 -m pip install nvidia-cublas==13.1.0.3 --force-reinstall --no-deps ; \
python3 -m pip install nixl-cu13 ; \
python3 -m pip install nvidia-nvshmem-cu13==${NVSHMEM_VERSION} --force-reinstall --no-deps; \
fi
# Install development tools and utilities
RUN --mount=type=cache,target=/var/cache/apt apt-get update && apt-get install -y \
gdb \