From f2f6d84e2ed0d25b3416e79709fd31e17e59adc8 Mon Sep 17 00:00:00 2001 From: Giuseppe Grossi Date: Wed, 2 Sep 2026 18:00:22 +0000 Subject: [PATCH] Add mooncake as pip install from rocm wheels Signed-off-by: Giuseppe Grossi --- .buildkite/scripts/ci-bake-rocm.sh | 43 +----------- docker/Dockerfile.rocm | 104 +---------------------------- docker/ci-rocm.hcl | 26 +------- requirements/rocm.txt | 2 + 4 files changed, 7 insertions(+), 168 deletions(-) diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index be736d0f6bff..15f8274fc62b 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm" DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl" DEFAULT_CI_BASE_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt tools/install_torchcodec_rocm.sh rust-toolchain.toml tests/vllm_test_utils" DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm" -DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust-toolchain-input rust-toolchain build_nixl lmcache_source build_lmcache build_rocshmem build_deepep mooncake_source build_mooncake package_mooncake export_mooncake mori_base ci_base" +DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust-toolchain-input rust-toolchain build_nixl lmcache_source build_lmcache build_rocshmem build_deepep mori_base ci_base" DEFAULT_CI_BASE_METADATA_VERSION="3" # ROCm CI forces REMOTE_VLLM=0, so content identity covers only the selected # local-source stages rather than unreachable remote-fetch alternatives. @@ -1684,12 +1684,9 @@ ci_base_metadata_pairs() { metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}" metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")" metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")" - metadata_pair "vllm.rocm.mooncake_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "MOONCAKE_REPO")" - metadata_pair "vllm.rocm.mooncake_commit" "${MOONCAKE_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "MOONCAKE_BRANCH")}" metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}" metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}" metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}" - metadata_pair "vllm.rocm.mooncake_cache_key" "${MOONCAKE_CACHE_KEY:-}" } write_ci_base_metadata_annotations() { @@ -2271,7 +2268,7 @@ extract_dependency_pins() { return 0 fi - for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH MOONCAKE_BRANCH; do + for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do if [[ -n "${!var:-}" ]]; then echo "Using provided ${var}: ${!var}" continue @@ -2295,11 +2292,9 @@ compute_dependency_cache_keys() { local ucx_branch="" local rocshmem_branch="" local deepep_branch="" - local mooncake_branch="" local nixl_material="" local rocshmem_material="" local deepep_material="" - local mooncake_material="" bake_dir=$(dirname "${VLLM_BAKE_FILE}") dockerfile_rocm="${bake_dir}/Dockerfile.rocm" @@ -2307,7 +2302,6 @@ compute_dependency_cache_keys() { ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH") rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH") deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH") - mooncake_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "MOONCAKE_BRANCH") if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl") @@ -2341,19 +2335,6 @@ compute_dependency_cache_keys() { export DEEPEP_CACHE_KEY echo "DeepEP dependency cache key: ${DEEPEP_CACHE_KEY}" fi - - if [[ -n "${mooncake_branch}" ]]; then - mooncake_material=$(compose_stage_cache_material \ - "${dockerfile_rocm}" \ - "base mooncake_source build_mooncake package_mooncake export_mooncake") - MOONCAKE_CACHE_KEY=$( - compose_dependency_cache_key \ - "${mooncake_branch}" \ - "${mooncake_material}" - ) - export MOONCAKE_CACHE_KEY - echo "Mooncake dependency cache key: ${MOONCAKE_CACHE_KEY}" - fi } compose_stage_cache_material() { @@ -2402,13 +2383,6 @@ dependency_cache_ref_for_target() { printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH:-}" fi ;; - mooncake-rocm-ci) - if [[ -n "${MOONCAKE_CACHE_KEY:-}" ]]; then - printf '%s\n' "${cache_repo}:mooncake-rocm-${MOONCAKE_CACHE_KEY}" - elif [[ -n "${MOONCAKE_BRANCH:-}" ]]; then - printf '%s\n' "${cache_repo}:mooncake-rocm-${MOONCAKE_BRANCH}" - fi - ;; esac } @@ -2426,14 +2400,13 @@ resolve_ci_base_dependency_targets() { local nixl_ref="" local rocshmem_ref="" local deepep_ref="" - local mooncake_ref="" [[ "${TARGET}" == "ci-base-rocm-ci-with-deps" ]] || return 0 case "${mode}" in always) echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially" - for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci mooncake-rocm-ci; do + for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then add_dependency_cache_target "${target}" fi @@ -2483,16 +2456,6 @@ resolve_ci_base_dependency_targets() { fi fi - if [[ "${mode}" != "always" && -n "${MOONCAKE_CACHE_KEY:-}" ]]; then - mooncake_ref=$(dependency_cache_ref_for_target "mooncake-rocm-ci") - if dependency_cache_ref_exists "${mooncake_ref}"; then - echo "Mooncake dependency cache exists: ${mooncake_ref}" - else - echo "Mooncake dependency cache missing; will seed: ${mooncake_ref}" - add_dependency_cache_target "mooncake-rocm-ci" - fi - fi - # DeepEP inherits from ROCShmem. If ROCShmem is being seeded, seed DeepEP too # so the pair stays consistent for future ci_base rebuilds. if printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" | grep -qx "rocshmem-rocm-ci" \ diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index fd8bfe4eb37e..e0b466da7b55 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -7,10 +7,6 @@ ARG BASE_IMAGE=rocm/vllm-dev:base ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base ARG ROCM_TRITON_KERNELS_COMMIT=0f380657dbf3ee86eb57558ff71df24f03b5d4e7 ARG INSTALL_LMCACHE=false -ARG MOONCAKE_BRANCH=df508641d76be20803349bf715d3bfc5205a7b58 -ARG MOONCAKE_REPO=https://github.com/kvcache-ai/Mooncake.git -ARG MOONCAKE_GO_VERSION=1.25.10 -ARG MOONCAKE_GO_SHA256=42d4f7a32316aa66591eca7e89867256057a4264451aca10570a715b3637ba70 # NIC backend for MoRI RDMA support. # By default (all), drivers and userspace libraries for all supported NIC types # (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime. @@ -647,82 +643,6 @@ RUN /bin/bash -lc 'set -euo pipefail; \ *) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \ esac' -FROM base AS mooncake_source -ARG MOONCAKE_GO_VERSION -ARG MOONCAKE_GO_SHA256 -ENV PATH=/usr/local/go/bin:${PATH} -ENV GOTOOLCHAIN=local -RUN test "$(uname -m)" = "x86_64" \ - && test ! -e /usr/local/go \ - && curl --fail --location --retry 5 \ - --output /tmp/go.tar.gz \ - "https://go.dev/dl/go${MOONCAKE_GO_VERSION}.linux-amd64.tar.gz" \ - && echo "${MOONCAKE_GO_SHA256} /tmp/go.tar.gz" | sha256sum --check - \ - && tar -C /usr/local -xzf /tmp/go.tar.gz \ - && rm /tmp/go.tar.gz \ - && test "$(go env GOVERSION)" = "go${MOONCAKE_GO_VERSION}" - -# Build a pinned ROCm Mooncake wheel outside the long-lived ROCm base image. -ARG MOONCAKE_BRANCH -ARG MOONCAKE_REPO -RUN git init mooncake \ - && cd mooncake \ - && git remote add origin "${MOONCAKE_REPO}" \ - && git fetch --depth=1 origin "${MOONCAKE_BRANCH}" \ - && git checkout --detach FETCH_HEAD \ - && test "$(git rev-parse HEAD)" = "${MOONCAKE_BRANCH}" \ - && grep -qx 'GOVER=1.25.9' dependencies.sh \ - && sed -i "s/^GOVER=.*/GOVER=${MOONCAKE_GO_VERSION}/" dependencies.sh \ - && bash dependencies.sh -y - -FROM mooncake_source AS build_mooncake -ARG USE_SCCACHE -ENV HIPCC_COMPILE_FLAGS_APPEND=--gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/11 -RUN --mount=type=cache,id=vllm-rocm-mooncake-go-mod,target=/root/go/pkg/mod \ - --mount=type=cache,id=vllm-rocm-mooncake-go-build,target=/root/.cache/go-build \ - cd mooncake \ - && cmake -G Ninja -B build \ - -DUSE_HIP=ON -DUSE_CUDA=OFF -DWITH_EP=OFF \ - -DUSE_HTTP=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON \ - -DBUILD_UNIT_TESTS=OFF -DBUILD_BENCHMARK=OFF \ - -DENABLE_DEBUG_SYMBOLS=OFF \ - -DENABLE_SCCACHE="${USE_SCCACHE:-OFF}" \ - -DCMAKE_BUILD_TYPE=Release \ - -DPython3_EXECUTABLE="$(command -v python3)" \ - && cmake --build build -j"${MAX_JOBS:-$(nproc)}" \ - && cmake --install build - -FROM build_mooncake AS package_mooncake -# Pin both the outer packaging tools and the isolated PEP 517 build environment. -# patchelf 0.14.3 from dependencies.sh corrupts auditwheel RPATHs. -ENV PIP_CONSTRAINT=/tmp/mooncake-build-constraints.txt -ENV PIP_BUILD_CONSTRAINT=/tmp/mooncake-build-constraints.txt -RUN --mount=type=cache,id=vllm-rocm-mooncake-pip,target=/root/.cache/pip \ - printf '%s\n' \ - 'pip==26.2.1' \ - 'build==1.5.0' \ - 'setuptools==79.0.1' \ - 'wheel==0.48.0' \ - 'auditwheel==6.8.1' \ - 'packaging==26.3' \ - 'pyproject-hooks==1.2.0' \ - 'pyelftools==0.33' \ - 'patchelf==0.19.1.0' \ - > "${PIP_CONSTRAINT}" \ - && python3 -m pip install --upgrade \ - pip build setuptools wheel auditwheel patchelf \ - && cd mooncake \ - && export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:/usr/local/lib" \ - && HIP_BUILD=1 OUTPUT_DIR=dist PLATFORM_TAG=manylinux_2_35_x86_64 \ - ./scripts/build_wheel.sh \ - && test "$(find mooncake-wheel/dist -maxdepth 1 -type f \ - -name 'mooncake_transfer_engine_rocm-*.whl' | wc -l)" -eq 1 \ - && mkdir -p /app/install \ - && cp mooncake-wheel/dist/mooncake_transfer_engine_rocm-*.whl /app/install/ - -FROM scratch AS export_mooncake -COPY --from=package_mooncake /app/install/ / - # ----------------------- # vLLM wheel release build stage (for building distributable wheels) # This stage pins dependencies to custom ROCm wheel versions and handles version detection @@ -953,15 +873,6 @@ RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_instal /lmcache_install/*.whl \ && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)" -RUN --mount=type=bind,from=export_mooncake,src=/,target=/mooncake_install,ro \ - --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system --no-deps \ - /mooncake_install/mooncake_transfer_engine_rocm-*.whl \ - && python3 -c "import aiohttp, msgpack, requests" \ - && python3 -c "from mooncake.engine import SUPPORT_HIP, TransferEngine; assert SUPPORT_HIP" \ - && python3 -c "from mooncake.store import MooncakeDistributedStore" \ - && mooncake_master --version - # ----------------------- # Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base. FROM ${CI_BASE_IMAGE} AS test @@ -1041,21 +952,10 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \ uv pip install --system /nixl_install/*.whl -RUN --mount=type=bind,from=export_mooncake,src=/,target=/mooncake_install,ro \ - --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system --no-deps \ - /mooncake_install/mooncake_transfer_engine_rocm-*.whl \ - && python3 -c "import aiohttp, msgpack, requests" \ - && python3 -c "from mooncake.engine import SUPPORT_HIP, TransferEngine; assert SUPPORT_HIP" \ - && python3 -c "from mooncake.store import MooncakeDistributedStore" \ - && mooncake_master --version - ARG COMMON_WORKDIR ARG BASE_IMAGE ARG NIC_BACKEND ARG AINIC_VERSION -ARG MOONCAKE_BRANCH -ARG MOONCAKE_REPO # Copy over the benchmark scripts as well COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks @@ -1084,9 +984,7 @@ RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \ && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \ - && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt \ - && echo "MOONCAKE_BRANCH=${MOONCAKE_BRANCH}" >> ${COMMON_WORKDIR}/versions.txt \ - && echo "MOONCAKE_REPO=${MOONCAKE_REPO}" >> ${COMMON_WORKDIR}/versions.txt + && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt CMD ["/bin/bash"] diff --git a/docker/ci-rocm.hcl b/docker/ci-rocm.hcl index 6a7053a356b4..1794b8468a16 100644 --- a/docker/ci-rocm.hcl +++ b/docker/ci-rocm.hcl @@ -91,10 +91,6 @@ variable "DEEPEP_BRANCH" { default = "" } -variable "MOONCAKE_BRANCH" { - default = "" -} - variable "NIXL_CACHE_KEY" { default = "" } @@ -107,10 +103,6 @@ variable "DEEPEP_CACHE_KEY" { default = "" } -variable "MOONCAKE_CACHE_KEY" { - default = "" -} - # Docker Hub registry cache for AMD builds. # # A separate repo (rocm/vllm-ci-cache) is used for BuildKit layer cache. @@ -255,7 +247,6 @@ function "get_cache_from_rocm_deps" { NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""), ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""), DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""), - MOONCAKE_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:mooncake-rocm-${MOONCAKE_CACHE_KEY}" : (MOONCAKE_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:mooncake-rocm-${MOONCAKE_BRANCH}" : ""), ]) } @@ -280,13 +271,6 @@ function "get_cache_to_rocm_deepep" { ]) } -function "get_cache_to_rocm_mooncake" { - params = [] - result = compact([ - MOONCAKE_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:mooncake-rocm-${MOONCAKE_CACHE_KEY},mode=min" : (MOONCAKE_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:mooncake-rocm-${MOONCAKE_BRANCH},mode=min" : ""), - ]) -} - # CI targets target "_ci-rocm" { @@ -417,14 +401,6 @@ target "deepep-rocm-ci" { output = ["type=cacheonly"] } -target "mooncake-rocm-ci" { - inherits = ["_common-rocm", "_ci-rocm"] - target = "export_mooncake" - cache-from = get_cache_from_rocm_deps() - cache-to = get_cache_to_rocm_mooncake() - output = ["type=cacheonly"] -} - # Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.) # Invoked by the ensure-ci-base step when the content hash of ci_base-affecting # files drifts from the remote image label. Per-PR builds then pull the result @@ -455,5 +431,5 @@ target "ci-base-rocm-ci" { # Group for ci_base builds -- exports dependency stage caches alongside the # ci_base image so future rebuilds can reuse them independently. group "ci-base-rocm-ci-with-deps" { - targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "mooncake-rocm-ci", "ci-base-rocm-ci"] + targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"] } diff --git a/requirements/rocm.txt b/requirements/rocm.txt index d2bf9898a8bd..824e2e12019f 100644 --- a/requirements/rocm.txt +++ b/requirements/rocm.txt @@ -27,3 +27,5 @@ tilelang==0.1.10 apache-tvm-ffi==0.1.10 # Required for faster safetensors model loading fastsafetensors >= 0.3.3 +# Mooncake transfer engine for KV cache offloading/transfer (ROCm build, published on PyPI) +mooncake-transfer-engine-rocm >= 0.3.13