[CI/Build] Always build with EP in CI (#922)

UNIDY2002 · web-flow · commit f89c207cfa39 · 2025-10-17T14:17:13.000+08:00
* Always build with EP in CI

* Add missing LIBRARY_PATH

* Do not build tests and examples for the final steps to save up disk space

* Remove the requirement of USE_CUDA=ON for transfer-engine (without EP)
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -171,6 +171,7 @@ jobs:
       matrix:
         python-version: ['3.10', '3.12']
     env:
+      BUILD_WITH_EP: "1"
       SCCACHE_GHA_ENABLED: "true"
 
     steps:
@@ -193,7 +194,8 @@ jobs:
         cuda: '12.8.1'
         linux-local-args: '["--toolkit"]'
         method: 'network'
-        sub-packages: '["nvcc"]'
+        sub-packages: '["nvcc", "nvrtc-dev"]'
+        non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]'
 
     - name: Run sccache-cache
       uses: mozilla-actions/sccache-action@v0.0.9
@@ -213,6 +215,7 @@ jobs:
       run: |
         sudo apt update -y
         sudo bash -x dependencies.sh -y
+        pip install torch==2.8.0
       shell: bash
 
     - name: Build transfer engine only
@@ -231,7 +234,7 @@ jobs:
       run: |
         mkdir build
         cd build
-        cmake .. -DUSE_ETCD=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
+        cmake .. -DUSE_ETCD=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_EP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
       shell: bash
       # TODO: lack USE_NVMEOF,USE_CUDA,USE_MNNVL
 
@@ -259,7 +262,8 @@ jobs:
     - name: Configure project
       run: |
         cd build
-        cmake .. -DUSE_HTTP=ON -DENABLE_SCCACHE=ON
+        rm -r */tests
+        cmake .. -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DUSE_HTTP=ON -DENABLE_SCCACHE=ON
       shell: bash
 
     - name: Build project
@@ -294,100 +298,6 @@ jobs:
         name: mooncake-wheel-ubuntu-py${{ steps.generate_tag_flags.outputs.python_version_tag }}
         path: mooncake-wheel/dist-py${{ steps.generate_tag_flags.outputs.python_version_tag }}/*.whl
 
-  build-with-ep:
-    runs-on: ubuntu-22.04
-    strategy:
-      matrix:
-        python-version: ['3.10', '3.12']
-    env:
-      BUILD_WITH_EP: "1"
-      SCCACHE_GHA_ENABLED: "true"
-      SCCACHE_CACHE_SIZE: "2G"
-
-    steps:
-    - uses: actions/checkout@v4
-
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v5
-      with:
-        python-version: ${{ matrix.python-version }}
-
-    - name: Free up disk space
-      run: |
-        sudo rm -rf /usr/share/dotnet
-        sudo rm -rf /opt/ghc
-        sudo rm -rf /opt/hostedtoolcache/CodeQL
-
-    - name: Install CUDA Toolkit
-      uses: Jimver/cuda-toolkit@v0.2.24
-      with:
-        cuda: '12.8.1'
-        linux-local-args: '["--toolkit"]'
-        method: 'network'
-        sub-packages: '["nvcc", "nvrtc-dev"]'
-        non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]'
-
-    - name: Run sccache-cache
-      uses: mozilla-actions/sccache-action@v0.0.9
-
-    - name: Configure sccache
-      uses: actions/github-script@v7
-      with:
-        script: |
-          core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
-          core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
-
-    - name: Run sccache stat for check
-      shell: bash
-      run: ${SCCACHE_PATH} --show-stats
-
-    - name: Install dependencies
-      run: |
-        sudo apt update -y
-        sudo bash -x dependencies.sh -y
-        pip install toml-cli  # for updating the version
-        pip install --no-cache-dir torch==2.8.0
-      shell: bash
-
-    - name: Build transfer engine with EP
-      run: |
-        mkdir build
-        cd build
-        export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
-        export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:LIBRARY_PATH
-        export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
-        cmake .. -DUSE_ETCD=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_EP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=ON -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
-        make -j
-        sudo make install
-      shell: bash
-
-    - name: Build nvlink_allocator.so
-      run: |
-        mkdir -p build/mooncake-transfer-engine/nvlink-allocator
-        cd mooncake-transfer-engine/nvlink-allocator
-        bash build.sh --ci-build ../../build/mooncake-transfer-engine/nvlink-allocator/
-      shell: bash
-
-    - name: Generate Python version tag
-      id: generate_tag_flags
-      run: |
-        echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
-      shell: bash
-
-    - name: Build Python wheel
-      run: |
-        BASE_VERSION=$(toml get --toml-path mooncake-wheel/pyproject.toml project.version | tr -d '"')
-        toml set --toml-path mooncake-wheel/pyproject.toml project.version "${BASE_VERSION}+ep"
-        # Build wheel with specific Python version
-        PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_flags.outputs.python_version_tag }} ./scripts/build_wheel.sh
-      shell: bash
-
-    - name: Upload Python wheel artifact
-      uses: actions/upload-artifact@v4
-      with:
-        name: mooncake-wheel-ubuntu-py${{ steps.generate_tag_flags.outputs.python_version_tag }}+ep
-        path: mooncake-wheel/dist-py${{ steps.generate_tag_flags.outputs.python_version_tag }}/*.whl
-
   build-docker:
     name: Build Docker Image
     runs-on: ubuntu-22.04
diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml
@@ -15,93 +15,6 @@ jobs:
     strategy:
       matrix:
         python-version: ['3.8', '3.9', '3.10', '3.11', '3.12', '3.13']
-    steps:
-      - name: Checkout source
-        uses: actions/checkout@v4
-
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python-version }}
-
-      - name: Free up disk space
-        run: |
-          sudo rm -rf /usr/share/dotnet
-          sudo rm -rf /opt/ghc
-          sudo rm -rf /opt/hostedtoolcache/CodeQL
-
-      - name: Install CUDA Toolkit
-        uses: Jimver/cuda-toolkit@v0.2.24
-        with:
-          cuda: '12.8.1'
-          linux-local-args: '["--toolkit"]'
-          method: 'network'
-          sub-packages: '["nvcc"]'
-
-      - name: Run sccache-cache
-        uses: mozilla-actions/sccache-action@v0.0.9
-
-      - name: Configure sccache
-        uses: actions/github-script@v7
-        with:
-          script: |
-            core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || '');
-            core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || '');
-
-      - name: Run sccache stat for check
-        shell: bash
-        run: ${SCCACHE_PATH} --show-stats
-
-      - name: Configure project
-        run: |
-          sudo apt update -y
-          sudo bash -x dependencies.sh -y
-          mkdir build
-          cd build
-          cmake .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
-        shell: bash
-
-      - name: Build project
-        run: |
-          cd build
-          make -j
-          sudo make install
-        shell: bash
-
-      - name: Build nvlink_allocator.so
-        run: |
-          mkdir -p build/mooncake-transfer-engine/nvlink-allocator
-          cd mooncake-transfer-engine/nvlink-allocator
-          bash build.sh --ci-build ../../build/mooncake-transfer-engine/nvlink-allocator/
-        shell: bash
-
-      - name: Generate Python version tag
-        id: generate_tag_release
-        run: |
-          echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT
-        shell: bash
-
-      - name: Build Python wheel
-        run: |
-          # Set LD_LIBRARY_PATH for wheel building
-          export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
-          PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh
-        env:
-          VERSION: ${{ env.VERSION }}
-
-      - name: Upload Python wheel artifact
-        uses: actions/upload-artifact@v4
-        with:
-          name: mooncake-wheel-py${{ steps.generate_tag_release.outputs.python_version_tag }}
-          path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
-
-  build-with-ep:
-    runs-on: ubuntu-22.04
-    permissions:
-      contents: write
-    strategy:
-      matrix:
-        python-version: ['3.10', '3.12']
     env:
       BUILD_WITH_EP: "1"
     steps:
@@ -146,11 +59,10 @@ jobs:
         run: |
           sudo apt update -y
           sudo bash -x dependencies.sh -y
-          pip install toml-cli  # for updating the version
           pip install torch==2.8.0
           mkdir build
           cd build
-          cmake .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
+          cmake .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DWITH_EP=ON -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
         shell: bash
 
       - name: Build project
@@ -175,8 +87,6 @@ jobs:
 
       - name: Build Python wheel
         run: |
-          BASE_VERSION=$(toml get --toml-path mooncake-wheel/pyproject.toml project.version | tr -d '"')
-          toml set --toml-path mooncake-wheel/pyproject.toml project.version "${BASE_VERSION}+ep"
           # Set LD_LIBRARY_PATH for wheel building
           export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
           PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh
@@ -186,13 +96,11 @@ jobs:
       - name: Upload Python wheel artifact
         uses: actions/upload-artifact@v4
         with:
-          name: mooncake-wheel-py${{ steps.generate_tag_release.outputs.python_version_tag }}+ep
+          name: mooncake-wheel-py${{ steps.generate_tag_release.outputs.python_version_tag }}
           path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl
 
   publish-release:
-    needs:
-      - build
-      - build-with-ep
+    needs: build
     runs-on: ubuntu-22.04
     permissions:
       contents: write
diff --git a/doc/en/build.md b/doc/en/build.md
@@ -6,11 +6,6 @@ This document describes how to build Mooncake from source.
    ```bash
    pip3 install mooncake-transfer-engine --upgrade
    ```
-- To install with the Mooncake Backend and Mooncake EP support, use the following command:
-   ```bash
-   # replace torch2.8.0 with the corresponding version
-   pip3 install mooncake-transfer-engine==0.3.7+ep --upgrade
-   ```
 
 ## Automatic
 
diff --git a/mooncake-ep/include/mooncake_backend.h b/mooncake-ep/include/mooncake_backend.h
@@ -8,6 +8,8 @@
 
 namespace mooncake {
 
+std::string getCudaTopologyJson(const std::vector<std::string>& filter);
+
 class MooncakeBackend final : public ::c10d::Backend {
    public:
     struct MooncakeBackendOptions final : ::c10d::Backend::Options {
@@ -61,6 +63,7 @@ class MooncakeBackend final : public ::c10d::Backend {
     static void setHostIp(const std::string& hostIp) { hostIp_ = hostIp; }
 
     static void setDeviceFilter(std::vector<std::string> filters) {
+        hca_filters_ = filters;
         engine_.setWhitelistFilters(std::move(filters));
     }
 
@@ -72,6 +75,7 @@ class MooncakeBackend final : public ::c10d::Backend {
    private:
     static TransferEngine engine_;
     static Transport* transport_;
+    static std::vector<std::string> hca_filters_;
     static int backendIndex_;
     bool isCpu_{false};
     static std::string hostIp_;
diff --git a/mooncake-ep/src/CMakeLists.txt b/mooncake-ep/src/CMakeLists.txt
@@ -1,4 +1,4 @@
-add_library(mooncake_ep mooncake_backend.cpp mooncake_ep_buffer.cpp mooncake_ep_kernel.cu mooncake_worker.cu mooncake_worker_thread.cpp mooncake_ibgda/mlx5gda.cpp)
+add_library(mooncake_ep mooncake_backend.cpp mooncake_cuda_topology.cpp mooncake_ep_buffer.cpp mooncake_ep_kernel.cu mooncake_worker.cu mooncake_worker_thread.cpp mooncake_ibgda/mlx5gda.cpp)
 
 set_target_properties(mooncake_ep PROPERTIES POSITION_INDEPENDENT_CODE ON)
 target_link_libraries(mooncake_ep PUBLIC ${TORCH_LIBRARIES} transfer_engine ibverbs mlx5)
diff --git a/mooncake-ep/src/mooncake_backend.cpp b/mooncake-ep/src/mooncake_backend.cpp
@@ -16,8 +16,9 @@ constexpr const char* SPARSE_ERROR_MSG = "Sparse op not supported.";
 constexpr const char* REDUCE_DTYPE_ERROR_MSG = "Unsupported reduce dtype: ";
 
 std::string MooncakeBackend::hostIp_ = "127.0.0.1";
-TransferEngine MooncakeBackend::engine_ = TransferEngine(true);
+TransferEngine MooncakeBackend::engine_ = TransferEngine();
 Transport* MooncakeBackend::transport_ = nullptr;
+std::vector<std::string> MooncakeBackend::hca_filters_;
 int MooncakeBackend::backendIndex_ = 0;
 MooncakeWorker MooncakeBackend::worker_;
 
@@ -33,7 +34,11 @@ MooncakeBackend::MooncakeBackend(
     // Initialize transfer engine
     if (!transport_) {
         engine_.init(P2PHANDSHAKE, hostIp_);
-        transport_ = engine_.installTransport("rdma", nullptr);
+        std::string topology = getCudaTopologyJson(hca_filters_);
+        void** args = (void**)malloc(2 * sizeof(void*));
+        args[0] = (void*)topology.c_str();
+        args[1] = nullptr;
+        transport_ = engine_.installTransport("rdma", args);
         TORCH_CHECK(transport_ != nullptr,
                     c10::str("Failed to install transport"));
     }
diff --git a/mooncake-ep/src/mooncake_cuda_topology.cpp b/mooncake-ep/src/mooncake_cuda_topology.cpp
diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh
