[Build] Fix CUDA 12 KV connector dependency selection (#57945)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
Sep 22, 2026 at 21:13 UTC
d90f0eade5b56ca83ba03839591fe111c51790b3
3 files changed
+33
-7
.buildkite/scripts/install-kv-connectors.sh
+14
-3
@@ -9,9 +9,20 @@ if python3 -c "import torch; raise SystemExit(0 if torch.version.hip is not None
9
exit 0
10
fi
11
12
-REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
12
+# Select requirements before optional packages can change the installed Torch.
13
+CUDA_MAJOR=$(python3 -c 'import torch; print(torch.version.cuda.split(".", 1)[0])')
14
+DEFAULT_REQUIREMENTS=/vllm-workspace/requirements/kv_connectors.txt
15
+if [ "${CUDA_MAJOR}" = "12" ]; then
16
+ DEFAULT_REQUIREMENTS=/vllm-workspace/requirements/kv_connectors_cu12.txt
17
+fi
18
+REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-${DEFAULT_REQUIREMENTS}}"
19
14
-uv pip install --system -r "${REQUIREMENTS_FILE}"
20
+CONSTRAINTS_DIR=$(mktemp -d)
21
+trap 'rm -rf "${CONSTRAINTS_DIR}"' EXIT
22
+uv pip freeze --system > "${CONSTRAINTS_DIR}/installed.txt"
23
+grep -E '^(torch|torchaudio|torchvision|triton|cuda-bindings|cuda-python|cuda-toolkit)==' \
24
+ "${CONSTRAINTS_DIR}/installed.txt" > "${CONSTRAINTS_DIR}/constraints.txt"
25
+uv pip install --system -c "${CONSTRAINTS_DIR}/constraints.txt" -r "${REQUIREMENTS_FILE}"
26
27
KV_METADATA=$(python3 - <<'PY'
28
import importlib.metadata as metadata
@@ -56,5 +67,5 @@ PY
67
# build first to avoid a clash.
68
if [ "${CUDA_MAJOR}" = "13" ] && [ -n "${MOONCAKE_VERSION}" ]; then
69
uv pip uninstall --system mooncake-transfer-engine 2>/dev/null || true
59
- uv pip install --system "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"
70
+ uv pip install --system -c "${CONSTRAINTS_DIR}/constraints.txt" "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"
71
fi
docker/Dockerfile
+14
-4
@@ -1264,7 +1264,8 @@ ENV UV_HTTP_TIMEOUT=500
1264
ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0'
1265
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
1266
RUN --mount=type=cache,target=/opt/uv/cache \
1267
- --mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \
1267
+ --mount=type=bind,source=requirements,target=/tmp/vllm-requirements,ro \
1268
+ set -eu; \
1269
CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
1270
CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-'); \
1271
CUDA_HOME=/usr/local/cuda; \
@@ -1273,11 +1274,19 @@ RUN --mount=type=cache,target=/opt/uv/cache \
1274
libcublas-dev-${CUDA_VERSION_DASH} \
1275
libcusolver-dev-${CUDA_VERSION_DASH}"; \
1276
if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
1276
- uv pip install --system -r /tmp/kv_connectors.txt --no-build || ( \
1277
+ KV_REQUIREMENTS=/tmp/vllm-requirements/kv_connectors.txt; \
1278
+ if [ "${CUDA_MAJOR}" = "12" ]; then \
1279
+ KV_REQUIREMENTS=/tmp/vllm-requirements/kv_connectors_cu12.txt; \
1280
+ fi; \
1281
+ # Preserve the stack used to compile vLLM when resolving optional dependencies. \
1282
+ uv pip freeze --system > /tmp/kv-installed.txt; \
1283
+ grep -E '^(torch|torchaudio|torchvision|triton|cuda-bindings|cuda-python|cuda-toolkit)==' \
1284
+ /tmp/kv-installed.txt > /tmp/kv-constraints.txt; \
1285
+ uv pip install --system -c /tmp/kv-constraints.txt -r "${KV_REQUIREMENTS}" --no-build || ( \
1286
# if the above fails, install from source
1287
apt-get update -y && \
1288
apt-get install -y --no-install-recommends --allow-change-held-packages ${BUILD_PKGS} && \
1280
- uv pip install --system -r /tmp/kv_connectors.txt --no-build-isolation && \
1289
+ uv pip install --system -c /tmp/kv-constraints.txt -r "${KV_REQUIREMENTS}" --no-build-isolation && \
1290
apt-get purge -y ${BUILD_PKGS} && \
1291
# clean up -dev packages, keep runtime libraries
1292
rm -rf /var/lib/apt/lists/* \
@@ -1288,8 +1297,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \
1297
MOONCAKE_VERSION=$(python3 -c "import importlib.metadata as m; print(m.version('mooncake-transfer-engine'))" 2>/dev/null || true) && \
1298
if [ "${CUDA_MAJOR}" = "13" ] && [ -n "${MOONCAKE_VERSION}" ]; then \
1299
uv pip uninstall --system mooncake-transfer-engine 2>/dev/null || true; \
1291
- uv pip install --system "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"; \
1300
+ uv pip install --system -c /tmp/kv-constraints.txt "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"; \
1301
fi; \
1302
+ rm -f /tmp/kv-installed.txt /tmp/kv-constraints.txt; \
1303
fi
1304
1305
# Mooncake registers GPU memory for RDMA via nvidia-peermem by default. Hosts
requirements/kv_connectors_cu12.txt
new
+5
@@ -0,0 +1,5 @@
1
+# LMCache 0.5.x uses CUDA 13 dependencies; retain the CUDA 12 release here.
2
+lmcache == 0.4.4
3
+cupy-cuda12x != 14.1.0
4
+nixl == 1.4.1
5
+mooncake-transfer-engine >= 0.3.12