[EP] Support multiple torch versions (#1098)

* Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug * Debug
2025-11-28 14:25:45 +08:00 · 2025-11-28 14:25:45 +08:00 · 63f395e165
parent 0d8ba6f458
commit 63f395e165
9 changed files with 89 additions and 60 deletions
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@ -259,6 +259,8 @@ jobs:
        python-version: ['3.10', '3.12']
    env:
      BUILD_WITH_EP: "1"
+      EP_TORCH_VERSIONS: "2.8.0;2.9.0;2.9.1"
+      TORCH_CUDA_ARCH_LIST: "8.0;9.0"
      SCCACHE_GHA_ENABLED: "true"

    steps:
@ -304,7 +306,6 @@ jobs:
      run: |
        sudo apt update -y
        sudo bash -x dependencies.sh -y
-        pip install torch==2.8.0
        df -h
      shell: bash

--- a/.github/workflows/release-non-cuda.yaml
+++ b/.github/workflows/release-non-cuda.yaml
@ -51,7 +51,6 @@ jobs:
        run: |
          sudo apt update -y
          sudo bash -x dependencies.sh -y
-          pip install torch==2.8.0
          mkdir build
          cd build
          cmake .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=OFF -DWITH_EP=OFF -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
--- a/.github/workflows/release.yaml
+++ b/.github/workflows/release.yaml
@ -17,6 +17,8 @@ jobs:
        python-version: ['3.9', '3.10', '3.11', '3.12', '3.13']
    env:
      BUILD_WITH_EP: "1"
+      EP_TORCH_VERSIONS: "2.8.0;2.9.0;2.9.1"
+      TORCH_CUDA_ARCH_LIST: "8.0;9.0"
    steps:
      - name: Checkout source
        uses: actions/checkout@v4
@ -59,7 +61,6 @@ jobs:
        run: |
          sudo apt update -y
          sudo bash -x dependencies.sh -y
-          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
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@ -58,9 +58,7 @@ if (WITH_STORE)
 endif()

 if (WITH_EP)
-  message(STATUS "Mooncake EP will be built")
-  add_subdirectory(mooncake-ep)
-  include_directories(mooncake-ep/include)
+  message(WARNING "Option `WITH_EP` is deprecated. Mooncake EP now builds with setuptools. Please set environment variable BUILD_WITH_EP=1 to enable.")
 endif()

 add_subdirectory(mooncake-integration)
--- a/mooncake-ep/setup.py
+++ b/mooncake-ep/setup.py
@ -0,0 +1,46 @@
+import os
+import re
+
+from setuptools import setup
+import torch
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+
+torch_version = re.match(r"\d+(?:\.\d+)*", torch.__version__).group()
+version_suffix = "_" + torch_version.replace(".", "_")
+module_name = "mooncake.ep" + version_suffix
+
+abi_flag = int(torch._C._GLIBCXX_USE_CXX11_ABI)
+current_dir = os.path.abspath(os.path.dirname(__file__))
+
+
+setup(
+    name=module_name,
+    ext_modules=[
+        CUDAExtension(
+            name=module_name,
+            include_dirs=[
+                os.path.join(current_dir, "include"),
+                os.path.join(current_dir, "../mooncake-transfer-engine/include"),
+            ],
+            sources=[
+                "../mooncake-integration/ep/ep_py.cpp",
+                "src/mooncake_backend.cpp",
+                "src/mooncake_ep_buffer.cpp",
+                "src/mooncake_ep_kernel.cu",
+                "src/mooncake_worker.cu",
+                "src/mooncake_worker_thread.cpp",
+                "src/mooncake_ibgda/mlx5gda.cpp",
+            ],
+            extra_compile_args={
+                "cxx": [f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}", "-std=c++20", "-O3", "-g0"],
+                "nvcc": [f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}", "-std=c++20", "-Xcompiler", "-O3", "-Xcompiler", "-g0"],
+            },
+            libraries=["ibverbs", "mlx5"],
+            extra_objects=[
+                os.path.join(current_dir, "../mooncake-wheel/mooncake/engine.so"),
+            ],
+        )
+    ],
+    cmdclass={"build_ext": BuildExtension},
+)
--- a/mooncake-integration/ep/ep_py.cpp
+++ b/mooncake-integration/ep/ep_py.cpp
@ -32,11 +32,6 @@ c10::intrusive_ptr<c10d::Backend> createMooncakeCpuBackend(
 }

 __attribute__((constructor)) static void MooncakeBackendConstructor() {
-    auto version = py::module::import("torch")
-                       .attr("__version__")
-                       .attr("split")("+")
-                       .cast<std::vector<std::string>>()[0];
-    TORCH_CHECK(version == "2.8.0", "Mooncake Backend requires torch==2.8.0");
    py::object module = py::module::import("torch.distributed");
    py::object register_backend =
        module.attr("Backend").attr("register_backend");
@ -63,7 +58,7 @@ at::Tensor getActiveRanks(c10::intrusive_ptr<c10d::Backend> backend) {
    return mooncakeBackend->getActiveRanksTensor();
 }

-PYBIND11_MODULE(ep, m) {
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("createMooncakeBackend", &createMooncakeBackend);
    m.def("createMooncakeCpuBackend", &createMooncakeCpuBackend);
    m.def("set_host_ip", &MooncakeBackend::setHostIp);
--- a/mooncake-wheel/mooncake/ep.py
+++ b/mooncake-wheel/mooncake/ep.py
@ -0,0 +1,16 @@
+import importlib
+import re
+
+import torch
+
+torch_version = re.match(r"\d+(?:\.\d+)*", torch.__version__).group()
+version_suffix = "_" + torch_version.replace(".", "_")
+
+try:
+    backend_module = importlib.import_module("mooncake.ep" + version_suffix)
+except ModuleNotFoundError:
+    raise ImportError(
+        f"Mooncake EP was not built against torch=={torch_version}.\n"
+        f"Open an issue at https://github.com/kvcache-ai/Mooncake/issues."
+    )
+globals().update({k: v for k, v in backend_module.__dict__.items() if not k.startswith("_")})
--- a/mooncake-wheel/setup.py
+++ b/mooncake-wheel/setup.py
@ -1,4 +1,3 @@
-import os
 import sys
 import platform
 from setuptools import setup, Distribution
@ -109,40 +108,7 @@ class CustomBdistWheel(bdist_wheel):
 # ---------------------------------------------------------------------------
 # setup()
 # ---------------------------------------------------------------------------
-if int(os.getenv("BUILD_WITH_EP", "0")):
-    import torch
-    from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-    abi_flag = int(torch._C._GLIBCXX_USE_CXX11_ABI)
-    current_dir = os.path.abspath(os.path.dirname(__file__))
-    ext_modules = [
-        CUDAExtension(
-            name="mooncake.ep",
-            include_dirs=[
-                os.path.join(current_dir, "../mooncake-ep/include"),
-                os.path.join(current_dir, "../mooncake-transfer-engine/include"),
-            ],
-            sources=["../mooncake-integration/ep/ep_py.cpp"],
-            extra_compile_args={
-                "cxx": [f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}", "-std=c++20"],
-                "nvcc": [f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}", "-std=c++20"],
-            },
-            libraries=["ibverbs", "mlx5"],
-            extra_objects=[
-                os.path.join(current_dir, "../build/mooncake-ep/src/libmooncake_ep.a"),
-                os.path.join(current_dir, "mooncake/engine.so"),
-            ],
-        )
-    ]
-    setup(
-        distclass=BinaryDistribution,
-        cmdclass={
-            "bdist_wheel": CustomBdistWheel,
-            "build_ext": BuildExtension,
-        },
-        ext_modules=ext_modules,
-    )
-else:
-    setup(
-        distclass=BinaryDistribution,
-        cmdclass={"bdist_wheel": CustomBdistWheel},
-    )
+setup(
+    distclass=BinaryDistribution,
+    cmdclass={"bdist_wheel": CustomBdistWheel},
+)
--- a/scripts/build_wheel.sh
+++ b/scripts/build_wheel.sh
@ -64,6 +64,21 @@ else
    echo "Skipping libascend_transport_mem.so (not built - Ascend disabled)"
 fi

+if [ "$BUILD_WITH_EP" = "1" ]; then
+    echo "Building Mooncake EP"
+    cd mooncake-ep
+    if [ -z "$EP_TORCH_VERSIONS" ]; then
+        python setup.py build_ext --build-lib .
+    else
+        for version in ${EP_TORCH_VERSIONS//;/ }; do
+            pip install torch==$version
+            python setup.py build_ext --build-lib . --force  # Force build when torch version changes
+        done
+    fi
+    cp mooncake/*.so ../mooncake-wheel/mooncake/
+    cd ..
+fi
+
 echo "Building wheel package..."
 # Build the wheel package
 cd mooncake-wheel
@ -109,11 +124,7 @@ fi

 if [ "$PYTHON_VERSION" = "3.8" ]; then
    echo "Repairing wheel with auditwheel for platform: $PLATFORM_TAG"
-    if [ "$BUILD_WITH_EP" = "1" ]; then
-        python -m build --wheel --outdir ${OUTPUT_DIR} --no-isolation
-    else
-        python -m build --wheel --outdir ${OUTPUT_DIR}
-    fi
+    python -m build --wheel --outdir ${OUTPUT_DIR}

    echo "python 3.8 auditwheel does not support wild-cards..."
    PATTERNS=(
@ -217,11 +228,7 @@ if [ "$PYTHON_VERSION" = "3.8" ]; then
    auditwheel repair ${OUTPUT_DIR}/*.whl $EXCLUDE_OPTS -w ${REPAIRED_DIR}/ --plat ${PLATFORM_TAG}
 else
    echo "Repairing wheel with auditwheel for platform: $PLATFORM_TAG"
-    if [ "$BUILD_WITH_EP" = "1" ]; then
-        python -m build --wheel --outdir ${OUTPUT_DIR} --no-isolation
-    else
-        python -m build --wheel --outdir ${OUTPUT_DIR}
-    fi
+    python -m build --wheel --outdir ${OUTPUT_DIR}
    auditwheel repair ${OUTPUT_DIR}/*.whl \
    --exclude libcurl.so* \
    --exclude libibverbs.so* \