diff --git a/docker/Dockerfile b/docker/Dockerfile index 46c99c483..e282380b6 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -7,7 +7,7 @@ ARG BRANCH_TYPE=remote ARG GRACE_BLACKWELL=0 ARG HOPPER_SBO=0 -ARG GRACE_BLACKWELL_DEEPEP_BRANCH=gb200_blog_part_2 +ARG GRACE_BLACKWELL_DEEPEP_BRANCH=hybrid-ep ARG HOPPER_SBO_DEEPEP_COMMIT=9f2fc4b3182a51044ae7ecb6610f7c9c3258c4d6 ARG DEEPEP_COMMIT=9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee ARG BUILD_AND_DOWNLOAD_PARALLEL=8 @@ -270,9 +270,9 @@ WORKDIR /build # Clone DeepEP RUN set -eux; \ if [ "$GRACE_BLACKWELL" = "1" ]; then \ - git clone https://github.com/fzyzcjy/DeepEP.git && \ + git clone https://github.com/deepseek-ai/DeepEP.git -b ${GRACE_BLACKWELL_DEEPEP_BRANCH} && \ cd DeepEP && \ - git checkout ${GRACE_BLACKWELL_DEEPEP_BRANCH} && \ + 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 .. ; \ diff --git a/scripts/ci/cuda/ci_install_deepep.sh b/scripts/ci/cuda/ci_install_deepep.sh index c8a45e379..58a963ce2 100755 --- a/scripts/ci/cuda/ci_install_deepep.sh +++ b/scripts/ci/cuda/ci_install_deepep.sh @@ -99,11 +99,10 @@ apt-get install -y --no-install-recommends libfabric-dev || { DEEPEP_DIR=/root/.cache/deepep rm -rf ${DEEPEP_DIR} if [ "$GRACE_BLACKWELL" = "1" ]; then - # We use Tom's DeepEP fork for GB200 for now, which supports fp4 dispatch. - GRACE_BLACKWELL_DEEPEP_BRANCH=gb200_blog_part_2 - git clone https://github.com/fzyzcjy/DeepEP.git ${DEEPEP_DIR} && \ + GRACE_BLACKWELL_DEEPEP_BRANCH=hybrid-ep + git clone https://github.com/deepseek-ai/DeepEP.git -b ${GRACE_BLACKWELL_DEEPEP_BRANCH} ${DEEPEP_DIR} && \ pushd ${DEEPEP_DIR} && \ - git checkout ${GRACE_BLACKWELL_DEEPEP_BRANCH} && \ + git checkout d28bd676c2120573c9f1425f0c16c39faa4117e6 && \ sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh && \ popd else