[NVIDIA] Fix cutedsl backend of MoE (#12353)

scripts/ci/ci_install_deepep.sh

@@ -10,9 +10,20 @@ export LD_LIBRARY_PATH="${NVSHMEM_DIR}/lib:$LD_LIBRARY_PATH"
 export PATH="${NVSHMEM_DIR}/bin:$PATH"
 export CUDA_HOME=/usr/local/cuda
 
-if python3 -c "import deep_ep" >/dev/null 2>&1; then
-    echo "deep_ep is already installed or importable. Skipping installation."
-    exit 0
+GRACE_BLACKWELL=${GRACE_BLACKWELL:-0}
+# Detect architecture
+ARCH=$(uname -m)
+if [ "$ARCH" != "x86_64" ] && [ "$ARCH" != "aarch64" ]; then
+    echo "Unsupported architecture: $ARCH"
+    exit 1
+fi
+
+# It seems GB200 ci runner preinstalls some wrong version of deep_ep, so we cannot rely on it.
+if [ "$GRACE_BLACKWELL" != "1" ]; then
+    if python3 -c "import deep_ep" >/dev/null 2>&1; then
+        echo "deep_ep is already installed or importable. Skipping installation."
+        exit 0
+    fi
 fi
 
 # Install system dependencies
@@ -35,8 +46,10 @@ dpkg -i libgdrapi_*.deb
 dpkg -i gdrcopy-tests_*.deb
 dpkg -i gdrcopy_*.deb
 
-if [ ! -e "/usr/lib/x86_64-linux-gnu/libmlx5.so" ]; then
-    ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so
+# Set up library paths based on architecture
+LIB_PATH="/usr/lib/$ARCH-linux-gnu"
+if [ ! -e "$LIB_PATH/libmlx5.so" ]; then
+    ln -s $LIB_PATH/libmlx5.so.1 $LIB_PATH/libmlx5.so
 fi
 apt-get update && apt-get install -y libfabric-dev
 
@@ -45,6 +58,11 @@ cd /opt/nvshmem
 wget https://developer.download.nvidia.com/compute/redist/nvshmem/3.4.5/source/nvshmem_src_cuda12-all-all-3.4.5.tar.gz
 tar -xf nvshmem_src_cuda12-all-all-3.4.5.tar.gz
 mv nvshmem_src nvshmem && cd nvshmem
+if [ "$GRACE_BLACKWELL" = "1" ]; then
+    CUDA_ARCH="100;120"
+else
+    CUDA_ARCH="90"
+fi
 NVSHMEM_SHMEM_SUPPORT=0 \
 NVSHMEM_UCX_SUPPORT=0 \
 NVSHMEM_USE_NCCL=0 \
@@ -53,13 +71,45 @@ NVSHMEM_IBGDA_SUPPORT=1 \
 NVSHMEM_PMIX_SUPPORT=0 \
 NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \
 NVSHMEM_USE_GDRCOPY=1 \
-cmake -S . -B build/ -DCMAKE_INSTALL_PREFIX=/opt/nvshmem/install -DCMAKE_CUDA_ARCHITECTURES=90
+cmake -S . -B build/ -DCMAKE_INSTALL_PREFIX=/opt/nvshmem/install -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCH}
 cd build
 make -j$(nproc) install
 
 # Install DeepEP
-rm -rf /root/.cache/deepep && git clone https://github.com/deepseek-ai/DeepEP.git /root/.cache/deepep && cd /root/.cache/deepep && git checkout 9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee
-cd /root/.cache/deepep && python3 setup.py install
+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} && \
+    pushd ${DEEPEP_DIR} && \
+    git checkout ${GRACE_BLACKWELL_DEEPEP_BRANCH} && \
+    sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh && \
+    popd
+else
+    git clone https://github.com/deepseek-ai/DeepEP.git ${DEEPEP_DIR} && \
+    pushd ${DEEPEP_DIR} && \
+    git checkout 9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee && \
+    popd
+fi
+
+cd ${DEEPEP_DIR}
+if [ "$GRACE_BLACKWELL" = "1" ]; then
+    CUDA_VERSION=$(nvidia-smi | grep "CUDA Version" | head -n1 | awk '{print $9}')
+    if [ "$CUDA_VERSION" = "12.8" ]; then
+        CHOSEN_TORCH_CUDA_ARCH_LIST='10.0'
+    elif awk -v ver="$CUDA_VERSION" 'BEGIN {exit !(ver > 12.8)}'; then
+        CHOSEN_TORCH_CUDA_ARCH_LIST='10.0;10.3'
+    else
+        echo "Unsupported CUDA version for Grace Blackwell: $CUDA_VERSION" && exit 1
+    fi && \
+    if [ "${CUDA_VERSION%%.*}" = "13" ]; then \
+        sed -i "/^    include_dirs = \['csrc\/'\]/a\    include_dirs.append('${CUDA_HOME}/include/cccl')" setup.py; \
+    fi
+    NVSHMEM_DIR=/opt/nvshmem/install TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" pip install --no-build-isolation .
+else
+    python3 setup.py install
+fi
 
 # Verify configuration
 echo "=== Verify NVSHMEM ==="