| 1 | #!/bin/bash |
| 2 | |
| 3 | set -ex |
| 4 | |
| 5 | # Default workspace directory |
| 6 | WORKSPACE=$(pwd)/eep_kernels_workspace |
| 7 | INSTALL_NVSHMEM=true |
| 8 | |
| 9 | # Parse command line arguments |
| 10 | while getopts "w:n" opt; do |
| 11 | case $opt in |
| 12 | w) |
| 13 | WORKSPACE="$OPTARG" |
| 14 | ;; |
| 15 | n) |
| 16 | INSTALL_NVSHMEM=false |
| 17 | ;; |
| 18 | \?) |
| 19 | echo "Invalid option: -$OPTARG" >&2 |
| 20 | exit 1 |
| 21 | ;; |
| 22 | esac |
| 23 | done |
| 24 | |
| 25 | if [ ! -d "$WORKSPACE" ]; then |
| 26 | mkdir -p "$WORKSPACE" |
| 27 | fi |
| 28 | |
| 29 | |
| 30 | # install dependencies if not installed |
| 31 | pip3 install cmake torch ninja |
| 32 | |
| 33 | # build nvshmem |
| 34 | pushd "$WORKSPACE" |
| 35 | # Reset NVSHMEM build if requested |
| 36 | if [ "$INSTALL_NVSHMEM" = true ]; then |
| 37 | mkdir -p nvshmem_src |
| 38 | wget https://developer.download.nvidia.com/compute/redist/nvshmem/3.2.5/source/nvshmem_src_3.2.5-1.txz |
| 39 | tar -xvf nvshmem_src_3.2.5-1.txz -C nvshmem_src --strip-components=1 |
| 40 | pushd nvshmem_src |
| 41 | wget https://github.com/deepseek-ai/DeepEP/raw/main/third-party/nvshmem.patch |
| 42 | git init |
| 43 | git apply -vvv nvshmem.patch |
| 44 | git apply --reject --whitespace=fix ../../eep_nvshmem.patch |
| 45 | else |
| 46 | pushd nvshmem_src |
| 47 | fi |
| 48 | |
| 49 | # assume CUDA_HOME is set correctly |
| 50 | if [ -z "$CUDA_HOME" ]; then |
| 51 | echo "CUDA_HOME is not set, please set it to your CUDA installation directory." |
| 52 | exit 1 |
| 53 | fi |
| 54 | |
| 55 | # disable all features except IBGDA |
| 56 | export NVSHMEM_IBGDA_SUPPORT=1 |
| 57 | |
| 58 | export NVSHMEM_SHMEM_SUPPORT=0 |
| 59 | export NVSHMEM_UCX_SUPPORT=0 |
| 60 | export NVSHMEM_USE_NCCL=0 |
| 61 | export NVSHMEM_PMIX_SUPPORT=0 |
| 62 | export NVSHMEM_TIMEOUT_DEVICE_POLLING=0 |
| 63 | export NVSHMEM_USE_GDRCOPY=0 |
| 64 | export NVSHMEM_IBRC_SUPPORT=0 |
| 65 | export NVSHMEM_BUILD_TESTS=0 |
| 66 | export NVSHMEM_BUILD_EXAMPLES=0 |
| 67 | export NVSHMEM_MPI_SUPPORT=0 |
| 68 | export NVSHMEM_BUILD_HYDRA_LAUNCHER=0 |
| 69 | export NVSHMEM_BUILD_TXZ_PACKAGE=0 |
| 70 | export NVSHMEM_TIMEOUT_DEVICE_POLLING=0 |
| 71 | |
| 72 | cmake -G Ninja -S . -B "$WORKSPACE"/nvshmem_build/ -DCMAKE_INSTALL_PREFIX="$WORKSPACE"/nvshmem_install |
| 73 | cmake --build "$WORKSPACE"/nvshmem_build/ --target install |
| 74 | |
| 75 | popd |
| 76 | |
| 77 | export CMAKE_PREFIX_PATH=$WORKSPACE/nvshmem_install:$CMAKE_PREFIX_PATH |
| 78 | |
| 79 |