[Build] Fix CUDA 12 KV connector dependency selection (#57945)

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

.buildkite/scripts/install-kv-connectors.sh

@@ -9,9 +9,20 @@ if python3 -c "import torch; raise SystemExit(0 if torch.version.hip is not None
     exit 0
 fi
 
-REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
+# Select requirements before optional packages can change the installed Torch.
+CUDA_MAJOR=$(python3 -c 'import torch; print(torch.version.cuda.split(".", 1)[0])')
+DEFAULT_REQUIREMENTS=/vllm-workspace/requirements/kv_connectors.txt
+if [ "${CUDA_MAJOR}" = "12" ]; then
+    DEFAULT_REQUIREMENTS=/vllm-workspace/requirements/kv_connectors_cu12.txt
+fi
+REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-${DEFAULT_REQUIREMENTS}}"
 
-uv pip install --system -r "${REQUIREMENTS_FILE}"
+CONSTRAINTS_DIR=$(mktemp -d)
+trap 'rm -rf "${CONSTRAINTS_DIR}"' EXIT
+uv pip freeze --system > "${CONSTRAINTS_DIR}/installed.txt"
+grep -E '^(torch|torchaudio|torchvision|triton|cuda-bindings|cuda-python|cuda-toolkit)==' \
+    "${CONSTRAINTS_DIR}/installed.txt" > "${CONSTRAINTS_DIR}/constraints.txt"
+uv pip install --system -c "${CONSTRAINTS_DIR}/constraints.txt" -r "${REQUIREMENTS_FILE}"
 
 KV_METADATA=$(python3 - <<'PY'
 import importlib.metadata as metadata
@@ -56,5 +67,5 @@ PY
 # build first to avoid a clash.
 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}"
+    uv pip install --system -c "${CONSTRAINTS_DIR}/constraints.txt" "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"
 fi

docker/Dockerfile

@@ -1264,7 +1264,8 @@ ENV UV_HTTP_TIMEOUT=500
 ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0'
 ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
 RUN --mount=type=cache,target=/opt/uv/cache \
-    --mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \
+    --mount=type=bind,source=requirements,target=/tmp/vllm-requirements,ro \
+    set -eu; \
     CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
     CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-'); \
     CUDA_HOME=/usr/local/cuda; \
@@ -1273,11 +1274,19 @@ RUN --mount=type=cache,target=/opt/uv/cache \
                 libcublas-dev-${CUDA_VERSION_DASH} \
                 libcusolver-dev-${CUDA_VERSION_DASH}"; \
     if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
-        uv pip install --system -r /tmp/kv_connectors.txt --no-build || ( \
+        KV_REQUIREMENTS=/tmp/vllm-requirements/kv_connectors.txt; \
+        if [ "${CUDA_MAJOR}" = "12" ]; then \
+            KV_REQUIREMENTS=/tmp/vllm-requirements/kv_connectors_cu12.txt; \
+        fi; \
+        # Preserve the stack used to compile vLLM when resolving optional dependencies. \
+        uv pip freeze --system > /tmp/kv-installed.txt; \
+        grep -E '^(torch|torchaudio|torchvision|triton|cuda-bindings|cuda-python|cuda-toolkit)==' \
+            /tmp/kv-installed.txt > /tmp/kv-constraints.txt; \
+        uv pip install --system -c /tmp/kv-constraints.txt -r "${KV_REQUIREMENTS}" --no-build || ( \
             # if the above fails, install from source
             apt-get update -y && \
             apt-get install -y --no-install-recommends --allow-change-held-packages ${BUILD_PKGS} && \
-            uv pip install --system -r /tmp/kv_connectors.txt --no-build-isolation && \
+            uv pip install --system -c /tmp/kv-constraints.txt -r "${KV_REQUIREMENTS}" --no-build-isolation && \
             apt-get purge -y ${BUILD_PKGS} && \
             # clean up -dev packages, keep runtime libraries
             rm -rf /var/lib/apt/lists/* \
@@ -1288,8 +1297,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \
         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}"; \
+            uv pip install --system -c /tmp/kv-constraints.txt "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"; \
         fi; \
+        rm -f /tmp/kv-installed.txt /tmp/kv-constraints.txt; \
     fi
 
 # Mooncake registers GPU memory for RDMA via nvidia-peermem by default. Hosts

requirements/kv_connectors_cu12.txt

@@ -0,0 +1,5 @@
+# LMCache 0.5.x uses CUDA 13 dependencies; retain the CUDA 12 release here.
+lmcache == 0.4.4
+cupy-cuda12x != 14.1.0
+nixl == 1.4.1
+mooncake-transfer-engine >= 0.3.12