[Script][Docker] add script for deepep

docker/Dockerfile.rocm_deepep_base

@@ -0,0 +1,87 @@
+# build image for vllm deepep rocm
+ARG AITER_BRANCH="wjx/grouped_gemm"
+ARG VLLM_BRANCH="ganyi/llama_fp8_03122025/deepep"
+ARG DEEPEP_BRANCH="develop-ali"
+ARG ARG_PYTORCH_ROCM_ARCH="gfx90a;gfx942"
+ARG COMMON_WORKDIR=/opt
+# user this image as default
+# 1. source build nvshmem and deepep
+# 2. reinstall aiter
+# 3. reinstall vllm
+ARG BASE_IMAGE=rocm/ali-private:ubuntu22.04_rocm6.4.3.127_aiter_6b586ae_vllm_5b842c2_20250911
+
+FROM ${BASE_IMAGE} AS base
+
+ARG ARG_PYTORCH_ROCM_ARCH
+ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
+
+# Install some basic utilities
+RUN apt-get update -q -y && apt-get install -q -y \
+    sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \
+    apt-transport-https ca-certificates wget curl libopenmpi-dev openmpi-bin libibverbs-dev \
+    && echo "✅ Successfully installed dependencies"
+
+ARG COMMON_WORKDIR
+COPY ./deepep_required/nvshmem-main/ ${COMMON_WORKDIR}/nvshmem-main
+COPY ./deepep_required/deepEP_private ${COMMON_WORKDIR}/deepEP_private
+COPY ./deepep_required/install_deepep.sh ${COMMON_WORKDIR}/install_deepep.sh
+RUN echo "✅ Successfully copied nv-shmem & deepEP_private source code"
+WORKDIR ${COMMON_WORKDIR}
+# source build nv-shmem
+ENV PATH=/opt/rocm-6.4.3/llvm/bin/:/opt/rocm-6.4.3/bin/:/usr/local/sbin/:/usr/local/bin/:/usr/sbin/:/usr/bin/
+RUN echo "✅ PATH is set to $PATH"
+RUN cd ./nvshmem-main/nvshmem_src_rocm_wip \
+    && cp runtime_patches/amd_warp_sync_functions.h /opt/rocm-6.4.3/include/hip/amd_detail/ \
+    && rm -rf build/ \
+    && echo "✅ SUCCESS: NV-SHMEM repository ready"
+
+# source build deepep
+RUN cd deepEP_private \
+    && git submodule sync \
+    && git submodule update --init --recursive \
+    && echo "✅ SUCCESS: DeepEP repository ready"
+# Replace with verification
+RUN if [ -f "/usr/local/lib/python3.10/dist-packages/torch/utils/cpp_extension.py" ]; then \
+        sed -i 's/-fno-gpu-rdc/-fgpu-rdc/g' /usr/local/lib/python3.10/dist-packages/torch/utils/cpp_extension.py && \
+        echo "Successfully replaced -fno-gpu-rdc with -fgpu-rdc"; \
+    else \
+        echo "Warning: cpp_extension.py not found"; \
+    fi
+ENV LD_LIBRARY_PATH=/opt/mpich/install/lib/:$LD_LIBRARY_PATH
+RUN echo "✅ LD_LIBRARY_PATH is set to $LD_LIBRARY_PATH"
+RUN cd deepEP_private && rm -rf build/
+
+ARG AITER_BRANCH
+# reinstall aiter
+RUN cd aiter \
+    && git fetch origin ${AITER_BRANCH} \
+    && git reset --hard \
+    && git checkout ${AITER_BRANCH} \
+    && python3 -m pip uninstall -y aiter \
+    && python3 setup.py develop \
+    && echo "✅ Successfully reinstalled Aiter"
+
+ARG VLLM_BRANCH
+# reinstall vllm
+RUN cd vllm \
+    && git fetch origin ${VLLM_BRANCH} \
+    && git reset --hard \
+    && git checkout ${VLLM_BRANCH} \
+    && python3 -m pip uninstall -y vllm \
+    && python3 setup.py develop \
+    && echo "✅ Successfully reinstalled VLLM"
+
+ENV RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1
+ENV RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1
+ENV TOKENIZERS_PARALLELISM=false
+
+# ENV that can improve safe tensor loading, and end-to-end time
+ENV SAFETENSORS_FAST_GPU=1
+
+# Performance environment variable.
+ENV HIP_FORCE_DEV_KERNARG=1
+
+# Enable Aiter. Make sure this only exists on the aiter branch.
+# ENV VLLM_ROCM_USE_AITER=1
+
+CMD ["/bin/bash"]

docker/deepep_required/README.md

@@ -0,0 +1,53 @@
+# Setup environment for deepep on ROCM
+This guide is based on https://amd.atlassian.net/wiki/spaces/MLSE/pages/1119057874/vLLM+deployment+with+DeepEP.
+## Prepare nv-shmem and deepep repo
+nv-shmem, https://github.com/anghostcici/nvshmem.git, is a private repo owned by Nicole.Nie. You may use this zip to acquire the latest nvshmem for its build [nvshmem-main.zip](https://amdcloud-my.sharepoint.com/:u:/g/personal/danli103_amd_com/EUi9AszhzDBNiZk0P_fF_C0BP2TW60WWaQD6dxyykzyq5w?e=PGNapB) and unzip it locally.
+In case you cannot access this zip file, the alternative is to scp it:
+```shell
+scp -r gyu@hjbog-srdc-35.amd.com:/home/gyu/nvshmem-main.zip ./
+```
+Then clone deepep:
+```shell
+git clone -b develop-ali https://github.com/ROCm/deepEP_private.git
+```
+**NOTE: Do remember to put nv-shmem and deepep under ./deepep_required/**
+## Build docker image
+```shell
+DOCKER_BUILDKIT=1 docker build --no-cache -f Dockerfile.rocm_deepep_base -t rocm/vllm_deepep .
+```
+## Launch docker and install deepep inside container
+```shell
+docker run -itd --rm --cap-add=SYS_PTRACE -e SHELL=/bin/bash --network=host --security-opt seccomp=unconfined --device=/dev/kfd --device=/dev/dri -v /home/gyu:/home/gyu/ -v /mnt:/mnt --group-add video --ipc=host --name ${CONTAINER_NAME} rocm/vllm_deepep
+
+docker exec -it ${CONTAINER_NAME} bash
+
+bash install_deepep.sh
+```
+## Run vLLM with deepep_high_throughput mode
+Currently only eager mode is verified.
+```shell
+# launch server
+export HF_HOME=/mnt/raid0/models/
+export VLLM_USE_V1=1
+export SAFETENSORS_FAST_GPU=1
+export VLLM_ROCM_USE_AITER=1
+export VLLM_ROCM_USE_AITER_MOE=1
+export VLLM_USE_TRITON_FLASH_ATTN=0
+export NCCL_DEBUG=WARN
+export VLLM_ALL2ALL_BACKEND=deepep_high_throughput
+export VLLM_LOGGING_LEVEL=DEBUG
+
+vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 \
+  --trust-remote-code \
+  --disable-log-requests \
+  --max-model-len 32768 \
+  --tensor-parallel-size 2 \
+  --data-parallel-size 2 \
+  --quantization fp8 \
+  --no-enable-prefix-caching \
+  --max_num_batched_tokens 32768 \
+  --enable-expert-parallel \
+  --enforce-eager \
+  --max_seq_len_to_capture 32768
+```
+

docker/deepep_required/install_deepep.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+set -ex
+
+# install nvshmem
+cd /opt/nvshmem-main/nvshmem_src_rocm_wip
+./configure && make -j -C build
+echo "✅ Successfully built NVSHMEM"
+
+# install deepep
+cd /opt/deepEP_private
+AITER_MOE=1 ROCM_HOME=/opt/rocm-6.4.3/ OMPI_DIR=/opt/mpich/install ROCSHMEM_DIR=/opt/nvshmem-main/nvshmem_src_rocm_wip/build/src/ python3 setup.py --variance rocm build develop --force-nvshmem-api
+echo "✅ Successfully built and installed DeepEP"