Docker: install DeepEP from release wheels (#34041)

2026-08-07 23:51:09 -07:00
parent c69d59395b
commit c9444deef4
4 changed files with 16 additions and 95 deletions
@@ -122,7 +122,6 @@ jobs:
            -f docker/Dockerfile \
            --build-arg CUDA_VERSION=12.9.2 \
            --build-arg BUILD_TYPE=all \
-            --build-arg GRACE_BLACKWELL=0 \
            --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \
            "${METADATA_ARGS[@]}" \
            ${VERSION_ARG} \
@@ -151,7 +150,6 @@ jobs:
            -f docker/Dockerfile \
            --build-arg CUDA_VERSION=13.0.3 \
            --build-arg BUILD_TYPE=all \
-            --build-arg GRACE_BLACKWELL=0 \
            --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \
            "${METADATA_ARGS[@]}" \
            ${VERSION_ARG} \
@@ -234,7 +232,6 @@ jobs:
            -f docker/Dockerfile \
            --build-arg CUDA_VERSION=12.9.2 \
            --build-arg BUILD_TYPE=all \
-            --build-arg GRACE_BLACKWELL=1 \
            --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \
            "${METADATA_ARGS[@]}" \
            ${VERSION_ARG} \
@@ -263,7 +260,6 @@ jobs:
            -f docker/Dockerfile \
            --build-arg CUDA_VERSION=13.0.3 \
            --build-arg BUILD_TYPE=all \
-            --build-arg GRACE_BLACKWELL=1 \
            --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \
            "${METADATA_ARGS[@]}" \
            ${VERSION_ARG} \
@@ -198,7 +198,6 @@ jobs:
            --build-arg CUDA_VERSION=13.0.3 \
            --build-arg BUILD_TYPE=all \
            --build-arg CMAKE_BUILD_PARALLEL_LEVEL=$(nproc) \
-            --build-arg GRACE_BLACKWELL=1 \
            --build-arg BRANCH_TYPE=local \
            --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \
            --no-cache \
@@ -178,7 +178,6 @@ jobs:
            -f docker/Dockerfile \
            --build-arg CUDA_VERSION=13.0.3 \
            --build-arg BUILD_TYPE=all \
-            --build-arg GRACE_BLACKWELL=0 \
            --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \
            ${{ needs.prepare.outputs.extra_build_args }} \
            --no-cache \
@@ -4,12 +4,6 @@ FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu24.04 AS base
 ARG TARGETARCH
 ARG BUILD_TYPE=all
 ARG BRANCH_TYPE=remote
-ARG GRACE_BLACKWELL=0
-ARG HOPPER_SBO=0
-
-ARG HOPPER_SBO_DEEPEP_COMMIT=9f2fc4b3182a51044ae7ecb6610f7c9c3258c4d6
-ARG DEEPEP_COMMIT=9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee
-ARG BUILD_AND_DOWNLOAD_PARALLEL=8
 ARG SGL_KERNEL_VERSION=0.4.6.post1
 ARG SGL_VERSION
 ARG SGL_DEEP_GEMM_VERSION=0.1.5.post2
@@ -165,8 +159,7 @@ ENV LANG=en_US.UTF-8 \
 #
 #   base
 #     |
-#     +-- torch_deps ------> deepep_builder (needs torch)
-#     |                  |-> flashinfer_cache (needs flashinfer)
+#     +-- torch_deps ------> flashinfer_cache (needs flashinfer)
 #     |                  \-> hpc_ops_builder (cmake-only build)
 #     |
 #     +-- devtools_builder (independent)
@@ -229,8 +222,10 @@ COPY rust/sglang-grpc      /tmp/sglang_deps/rust/sglang-grpc
 COPY rust/sglang-mm        /tmp/sglang_deps/rust/sglang-mm
 COPY proto                 /tmp/sglang_deps/proto

-# Install sglang dependencies (torch, transformers, etc.)
-# Generate constraints.txt to prevent reinstalling these deps in later stages
+# Install sglang dependencies (torch, transformers, etc.). CUDA 12 DeepEP
+# wheels live only on the SGLang index, so preinstall the local-version wheel;
+# it satisfies the public-version pyproject pin during the full dependency solve.
+# Generate constraints.txt to prevent reinstalling these deps in later stages.
 RUN --mount=type=cache,target=/root/.cache/pip \
    --mount=type=cache,target=/root/.cargo/registry \
    case "$CUDA_VERSION" in \
@@ -245,6 +240,14 @@ RUN --mount=type=cache,target=/root/.cache/pip \
    && echo '__version__ = "0.0.0"' > sglang/version.py \
    && touch README.md \
    && touch LICENSE \
+    && SGL_DEEP_EP_VERSION="$(sed -n 's/^[[:space:]]*"sgl-deep-ep==\([^"]*\)",/\1/p' pyproject.toml)" \
+    && test -n "${SGL_DEEP_EP_VERSION}" \
+    && if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
+           python3 -m pip install \
+               "sgl-deep-ep==${SGL_DEEP_EP_VERSION}+cu129" \
+               --index-url "https://docs.sglang.ai/whl/cu129/" \
+               --no-deps; \
+       fi \
    && if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
           sed -i 's/cuda-python>=13\.0/cuda-python>=12,<13/' pyproject.toml && \
           sed -i 's/flashinfer_python\[cu13\]/flashinfer_python[cu12]/' pyproject.toml && \
@@ -267,76 +270,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 RUN python3 -m pip install --ignore-installed --no-deps distro

 ########################################################
-# PARALLEL STAGE 2: DeepEP Builder (needs torch_deps)
-########################################################
-FROM torch_deps AS deepep_builder
-
-ARG CUDA_VERSION
-ARG BUILD_AND_DOWNLOAD_PARALLEL
-ARG GRACE_BLACKWELL
-ARG HOPPER_SBO
-ARG HOPPER_SBO_DEEPEP_COMMIT
-ARG DEEPEP_COMMIT
-ARG GITHUB_ARTIFACTORY
-
-WORKDIR /build
-
-# Clone DeepEP
-RUN set -eux; \
-    if [ "$GRACE_BLACKWELL" = "1" ]; then \
-      if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
-        git clone https://github.com/fzyzcjy/DeepEP.git && \
-        cd DeepEP && \
-        git checkout gb200_blog_part_2 && \
-        sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh && \
-        sed -i 's/#define NUM_TIMEOUT_CYCLES 200000000000ull/#define NUM_TIMEOUT_CYCLES 2000000000000ull/' csrc/kernels/configs.cuh && \
-        cd .. ; \
-      else \
-        git clone https://github.com/deepseek-ai/DeepEP.git -b hybrid-ep && \
-        cd DeepEP && \
-        git checkout d28bd676c2120573c9f1425f0c16c39faa4117e6 && \
-        sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh && \
-        sed -i 's/#define NUM_TIMEOUT_CYCLES 200000000000ull/#define NUM_TIMEOUT_CYCLES 2000000000000ull/' csrc/kernels/configs.cuh && \
-        cd .. ; \
-      fi; \
-    elif [ "$HOPPER_SBO" = "1" ]; then \
-      git clone https://github.com/deepseek-ai/DeepEP.git -b antgroup-opt && \
-      cd DeepEP && \
-      git checkout ${HOPPER_SBO_DEEPEP_COMMIT} && \
-      sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh && \
-      sed -i 's/#define NUM_TIMEOUT_CYCLES 200000000000ull/#define NUM_TIMEOUT_CYCLES 2000000000000ull/' csrc/kernels/configs.cuh && \
-      cd .. ; \
-    else \
-        curl --retry 3 --retry-delay 2 -fsSL -o ${DEEPEP_COMMIT}.zip \
-            https://${GITHUB_ARTIFACTORY}/deepseek-ai/DeepEP/archive/${DEEPEP_COMMIT}.zip && \
-        unzip -q ${DEEPEP_COMMIT}.zip && rm ${DEEPEP_COMMIT}.zip && mv DeepEP-${DEEPEP_COMMIT} DeepEP && cd DeepEP && \
-        sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh && \
-        sed -i 's/#define NUM_TIMEOUT_CYCLES 200000000000ull/#define NUM_TIMEOUT_CYCLES 2000000000000ull/' csrc/kernels/configs.cuh && \
-        cd .. ; \
-    fi
-
-# Build DeepEP wheel
-RUN --mount=type=cache,target=/root/.cache/pip \
-    cd /build/DeepEP && \
-    case "$CUDA_VERSION" in \
-        12.6.3) \
-            CHOSEN_TORCH_CUDA_ARCH_LIST='9.0' \
-            ;; \
-        12.9.2|13.0.3) \
-            CHOSEN_TORCH_CUDA_ARCH_LIST='9.0;10.0;10.3' \
-            ;; \
-        *) \
-            echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 \
-            ;; \
-    esac && \
-    if [ "${CUDA_VERSION%%.*}" = "13" ]; then \
-        sed -i "/^    include_dirs = \['csrc\/'\]/a\    include_dirs.append('${CUDA_HOME}/include/cccl')" setup.py; \
-    fi && \
-    TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" MAX_JOBS=${BUILD_AND_DOWNLOAD_PARALLEL} \
-        python3 setup.py bdist_wheel -d /wheels
-
-########################################################
-# PARALLEL STAGE 3: HPC-Ops Builder (needs torch_deps)
+# PARALLEL STAGE 2: HPC-Ops Builder (needs torch_deps)
 ########################################################
 FROM torch_deps AS hpc_ops_builder

@@ -360,7 +294,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
    fi

 ########################################################
-# PARALLEL STAGE 4: FlashInfer Cache (needs torch_deps)
+# PARALLEL STAGE 3: FlashInfer Cache (needs torch_deps)
 ########################################################
 FROM torch_deps AS flashinfer_cache

@@ -498,7 +432,6 @@ FROM torch_deps AS framework
 ARG BRANCH_TYPE
 ARG BUILD_TYPE
 ARG CUDA_VERSION
-ARG BUILD_AND_DOWNLOAD_PARALLEL
 ARG SGL_VERSION
 ARG USE_LATEST_SGLANG
 ARG GITHUB_ARTIFACTORY
@@ -512,12 +445,6 @@ WORKDIR /sgl-workspace
 # Copy artifacts from parallel builders
 # =============================================================================

-# Copy DeepEP wheel and install
-COPY --from=deepep_builder /wheels /tmp/wheels/deepep
-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 HPC-Ops wheel and install (empty on non-x86_64; kernels are sm90a-only)
 COPY --from=hpc_ops_builder /wheels /tmp/wheels/hpc-ops
 RUN --mount=type=cache,target=/root/.cache/pip \
@@ -788,7 +715,7 @@ LABEL org.opencontainers.image.source="https://github.com/sgl-project/sglang" \
 # PURPOSE: Production runtime environment with JIT support
 #
 # This stage creates a production-ready image containing:
-# - Pre-compiled SGLang and DeepEP components
+# - Pre-installed SGLang and CUDA dependencies
 # - Full CUDA toolchain for JIT compilation (DeepGEMM, Triton, FlashInfer)
 # - Optimized for inference workloads and deployment
 # - Smaller than framework (no dev tools like vim, tmux, nsight, etc.)