mirror of
https://github.com/vllm-project/vllm.git
synced 2026-08-16 18:58:06 +00:00
[Docker][KVConnector] Install mooncake from official wheels instead of a custom build (#51067)
Signed-off-by: Zhewen Li <[email protected]> Co-authored-by: Zhewen Li <[email protected]> Co-authored-by: Claude Opus 5 <[email protected]>
This commit is contained in:
co-authored by
Zhewen Li
Claude Opus 5
parent
613411a90c
commit
14e57ad47d
@@ -16,13 +16,6 @@ env:
|
||||
# please use CUDA 13 wheels or compile yourself on these new devices
|
||||
CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
|
||||
CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
|
||||
|
||||
# pre-built mooncake wheels
|
||||
# the manylinux_2_35 wheel has compatibility issue on Ubuntu 24.04
|
||||
# so we use different wheels for the time being
|
||||
MOONCAKE_WHEEL_AARCH64_2_35: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_aarch64.whl"
|
||||
MOONCAKE_WHEEL_AARCH64_2_39: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_39_aarch64.whl"
|
||||
MOONCAKE_WHEEL_X86_64: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_x86_64.whl"
|
||||
|
||||
steps:
|
||||
- input: "Provide Release version here"
|
||||
@@ -219,8 +212,6 @@ steps:
|
||||
--build-arg CUDA_VERSION=13.0.2 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
|
||||
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
|
||||
--build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
@@ -247,8 +238,6 @@ steps:
|
||||
--build-arg CUDA_VERSION=13.0.2 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
|
||||
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
|
||||
--build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
@@ -273,8 +262,6 @@ steps:
|
||||
--build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
|
||||
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
@@ -301,8 +288,6 @@ steps:
|
||||
--build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
|
||||
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
@@ -327,8 +312,6 @@ steps:
|
||||
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
|
||||
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
|
||||
--build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
@@ -356,8 +339,6 @@ steps:
|
||||
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
|
||||
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
|
||||
--build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
@@ -384,8 +365,6 @@ steps:
|
||||
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
|
||||
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
@@ -413,8 +392,6 @@ steps:
|
||||
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
|
||||
--build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
|
||||
+17
-23
@@ -975,32 +975,26 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
# clean up -dev packages, keep runtime libraries
|
||||
rm -rf /var/lib/apt/lists/* \
|
||||
); \
|
||||
# Force-reinstall the matching CUDA wheel so the correct nixl_ep_cpp.so is installed.
|
||||
uv pip install --system --force-reinstall --no-deps nixl-cu${CUDA_MAJOR}; \
|
||||
fi
|
||||
|
||||
# Optional override: install mooncake-transfer-engine from a URL instead of the
|
||||
# PyPI release pulled in above. Use this for wheels built with non-default CMake
|
||||
# flags (e.g. `STORE_USE_ETCD=ON` for master HA). The URL's manylinux glibc
|
||||
# floor must be <= the FINAL_BASE_IMAGE's glibc.
|
||||
ARG MOONCAKE_WHEEL_AARCH64
|
||||
ARG MOONCAKE_WHEEL_X86_64
|
||||
RUN if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
|
||||
if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
|
||||
WHEEL="${MOONCAKE_WHEEL_AARCH64}"; \
|
||||
else \
|
||||
WHEEL="${MOONCAKE_WHEEL_X86_64}"; \
|
||||
fi && \
|
||||
if [ -n "${WHEEL}" ]; then \
|
||||
uv pip install --system "${WHEEL}" && \
|
||||
CUDA_MAJOR="${CUDA_VERSION%%.*}" && \
|
||||
if [ ! -f /usr/local/cuda/lib64/libcudart.so ] && \
|
||||
[ -f "/usr/local/cuda/lib64/libcudart.so.${CUDA_MAJOR}" ]; then \
|
||||
ln -s "libcudart.so.${CUDA_MAJOR}" /usr/local/cuda/lib64/libcudart.so; \
|
||||
fi; \
|
||||
# Force-reinstall the matching CUDA wheel so the correct nixl_ep_cpp.so is
|
||||
# installed. Chained with && through the mooncake swap below: this RUN has no
|
||||
# `set -e`, so its status is that of the last command, and a trailing no-op
|
||||
# would otherwise mask a failure here.
|
||||
uv pip install --system --force-reinstall --no-deps nixl-cu${CUDA_MAJOR} && \
|
||||
# The default mooncake wheel links libcudart.so.12, so on CUDA 13 swap in the
|
||||
# cuda13 variant. Uninstall first; both ship the `mooncake` package.
|
||||
# Mirrors .buildkite/scripts/install-kv-connectors.sh.
|
||||
MOONCAKE_VERSION=$(python3 -c "import importlib.metadata as m; print(m.version('mooncake-transfer-engine'))" 2>/dev/null || true) && \
|
||||
if [ "${CUDA_MAJOR}" = "13" ] && [ -n "${MOONCAKE_VERSION}" ]; then \
|
||||
uv pip uninstall --system mooncake-transfer-engine 2>/dev/null || true; \
|
||||
uv pip install --system "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"; \
|
||||
fi; \
|
||||
fi
|
||||
|
||||
# Mooncake registers GPU memory for RDMA via nvidia-peermem by default. Hosts
|
||||
# without that module loaded (e.g. GB200) must pass WITH_NVIDIA_PEERMEM=0 at run
|
||||
# time; see docs/features/mooncake_connector_usage.md. Deliberately not an ENV
|
||||
# here so the image does not override mooncake's own default.
|
||||
|
||||
ENV VLLM_USAGE_SOURCE production-docker-image
|
||||
ENV VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT:-unknown} \
|
||||
VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE:-local} \
|
||||
|
||||
@@ -10,7 +10,9 @@ For more details about Mooncake, please refer to [Mooncake project](https://gith
|
||||
|
||||
### Installation
|
||||
|
||||
Install mooncake through pip: `uv pip install mooncake-transfer-engine`.
|
||||
Install mooncake through pip: `uv pip install mooncake-transfer-engine-cuda13`.
|
||||
|
||||
vLLM defaults to CUDA 13. On a CUDA 12 environment install `mooncake-transfer-engine` instead — the two are the same release built against different CUDA majors, and the wrong one fails to import with `libcudart.so.<major>: cannot open shared object file`.
|
||||
|
||||
Refer to [Mooncake official repository](https://github.com/kvcache-ai/Mooncake) for more installation instructions
|
||||
|
||||
@@ -44,6 +46,12 @@ Now you can send requests to the proxy server through port 8000.
|
||||
- For headless instances, must be the same as the master instance
|
||||
- Each instance needs a unique port on its host; using the same port number across different hosts is fine
|
||||
|
||||
- `WITH_NVIDIA_PEERMEM`: Selects how mooncake registers GPU memory for RDMA. Read by mooncake, not vLLM.
|
||||
- Default: 1, which uses `ibv_reg_mr()` and requires the `nvidia-peermem` kernel module to be loaded
|
||||
- Set to 0 to use the DMA-BUF path, which does not need that module. Required on hosts where `nvidia-peermem` is not loaded, such as GB200
|
||||
- With the container image, pass it at run time: `docker run -e WITH_NVIDIA_PEERMEM=0 ...`
|
||||
- Symptom when left unset on such a host: `Failed to register memory <addr>: Bad address [14]` from `rdma_context.cpp`, and KV transfers fail
|
||||
|
||||
- `VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT`: Timeout (in seconds) for automatically releasing the prefiller’s KV cache for a particular request. (Optional)
|
||||
- Default: 480
|
||||
- If a request is aborted and the decoder has not yet notified the prefiller, the prefill instance will release its KV-cache blocks after this timeout to avoid holding them indefinitely.
|
||||
|
||||
Reference in New Issue
Block a user