diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6f86c51a..ae10f123 100644 --- 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 diff --git a/.github/workflows/release-non-cuda.yaml b/.github/workflows/release-non-cuda.yaml index ad08c260..171a79bf 100644 --- 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 diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 02fec47b..7cdb10ad 100644 --- 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 diff --git a/CMakeLists.txt b/CMakeLists.txt index a664909d..ac708024 100644 --- 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) diff --git a/mooncake-ep/setup.py b/mooncake-ep/setup.py new file mode 100644 index 00000000..927a0113 --- /dev/null +++ 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}, +) diff --git a/mooncake-integration/ep/ep_py.cpp b/mooncake-integration/ep/ep_py.cpp index e9470e1c..07e409af 100644 --- a/mooncake-integration/ep/ep_py.cpp +++ b/mooncake-integration/ep/ep_py.cpp @@ -32,11 +32,6 @@ c10::intrusive_ptr createMooncakeCpuBackend( } __attribute__((constructor)) static void MooncakeBackendConstructor() { - auto version = py::module::import("torch") - .attr("__version__") - .attr("split")("+") - .cast>()[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 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); diff --git a/mooncake-wheel/mooncake/ep.py b/mooncake-wheel/mooncake/ep.py new file mode 100644 index 00000000..68edb799 --- /dev/null +++ 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("_")}) diff --git a/mooncake-wheel/setup.py b/mooncake-wheel/setup.py index ccb7b8ae..f0ead576 100644 --- 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}, +) diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh index 26861b7c..d94b8f26 100755 --- 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* \