fix(build): package native artifacts in CUDA wheels

.github/scripts/verify_cuda_wheel_artifacts.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Verify native artifacts required by a CUDA wheel."""
+
+from __future__ import annotations
+
+import argparse
+import stat
+import sys
+import zipfile
+from pathlib import Path
+
+import tomllib
+
+
+def supported_python_minors(pyproject: Path) -> list[int]:
+    with pyproject.open("rb") as file:
+        spec = tomllib.load(file)["project"]["requires-python"]
+    bounds = spec.removeprefix(">=3.").split(",<3.")
+    if len(bounds) != 2 or not all(bound.isdigit() for bound in bounds):
+        raise ValueError(f"Cannot parse requires-python from {pyproject}")
+    lower, upper = map(int, bounds)
+    return list(range(lower, upper))
+
+
+def verify(wheel: Path, pyproject: Path) -> None:
+    with zipfile.ZipFile(wheel) as archive:
+        entries = {entry.filename: entry for entry in archive.infolist()}
+
+    missing: list[str] = []
+    rust_binary = "aphrodite/aphrodite-rs"
+    if rust_binary not in entries:
+        missing.append(rust_binary)
+    else:
+        mode = entries[rust_binary].external_attr >> 16
+        if not mode & stat.S_IXUSR:
+            missing.append(f"{rust_binary} (not executable)")
+
+    if not any(
+        name.startswith("aphrodite/_rust_tool_parser")
+        and name.endswith(".so")
+        and "/" not in name.removeprefix("aphrodite/")
+        for name in entries
+    ):
+        missing.append("aphrodite/_rust_tool_parser*.so")
+
+    for minor in supported_python_minors(pyproject):
+        prefix = f"aphrodite/third_party/deep_gemm/_C.cpython-3{minor}-"
+        if not any(name.startswith(prefix) and name.endswith(".so") for name in entries):
+            missing.append(f"{prefix}*.so")
+
+    if missing:
+        formatted = "\n".join(f"  - {item}" for item in missing)
+        raise RuntimeError(f"{wheel} is missing required native artifacts:\n{formatted}")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("wheel", type=Path)
+    parser.add_argument("--pyproject", type=Path, default=Path("pyproject.toml"))
+    args = parser.parse_args()
+    verify(args.wheel, args.pyproject)
+    print(f"Verified CUDA wheel native artifacts: {args.wheel}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

docker/Dockerfile

@@ -227,6 +227,72 @@ ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
 ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
 #################### BUILD BASE IMAGE ####################
 
+#################### RUST BUILD IMAGE ####################
+# Build the Rust frontend separately so the CUDA wheel build does not need a
+# Rust toolchain and can reuse Cargo artifacts across wheel commits.
+FROM scratch AS rust-input
+COPY requirements/build/rust.txt /rust-input/aphrodite/requirements/build/rust.txt
+COPY rust/Cargo.lock rust/Cargo.toml /rust-input/aphrodite/rust/
+COPY rust/proto /rust-input/aphrodite/rust/proto
+COPY rust/src /rust-input/aphrodite/rust/src
+COPY tools/build_rust.py /rust-input/aphrodite/tools/build_rust.py
+
+FROM base AS rust-build
+ARG USE_SCCACHE
+ARG SCCACHE_ENDPOINT
+
+ENV CARGO_BUILD_JOBS=4
+ENV CARGO_NET_RETRY=10
+ENV RUSTUP_MAX_RETRIES=10
+ENV CARGO_HOME=/root/.cargo
+ENV RUSTUP_HOME=/root/.rustup
+ENV PATH=${CARGO_HOME}/bin:${PATH}
+ENV SCCACHE_SERVER_PORT=4227
+
+WORKDIR /workspace
+
+COPY --from=rust-input \
+    /rust-input/aphrodite/requirements/build/rust.txt \
+    requirements/build/rust.txt
+
+RUN apt-get update -y \
+    && apt-get install -y --no-install-recommends \
+        make \
+        unzip \
+        protobuf-compiler \
+        libprotobuf-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \
+      | sh -s -- -y --profile minimal --default-toolchain stable \
+    && rustc --version \
+    && cargo --version
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install \
+      --python /opt/venv/bin/python3 \
+      -r requirements/build/rust.txt
+
+COPY --from=rust-input /rust-input/aphrodite/rust rust
+COPY --from=rust-input \
+    /rust-input/aphrodite/tools/build_rust.py \
+    tools/build_rust.py
+
+RUN --mount=type=cache,id=aphrodite-cuda-cargo-registry,target=/root/.cargo/registry,sharing=locked \
+    --mount=type=cache,id=aphrodite-cuda-cargo-git,target=/root/.cargo/git,sharing=locked \
+    --mount=type=cache,id=aphrodite-cuda-cargo-target,target=/workspace/rust/target,sharing=locked \
+    --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \
+    if [ "$USE_SCCACHE" = "1" ]; then \
+        if [ -n "${SCCACHE_ENDPOINT}" ]; then export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; fi; \
+        export RUSTC_WRAPPER=sccache; \
+        sccache --show-stats; \
+    fi \
+    && python3 tools/build_rust.py --release \
+    && if [ "$USE_SCCACHE" = "1" ]; then \
+        sccache --show-stats; \
+    fi
+#################### RUST BUILD IMAGE ####################
+
 #################### EXTENSIONS BUILD IMAGE ####################
 # Build DeepEP wheels
 FROM base AS extensions-build
@@ -301,8 +367,18 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 
 WORKDIR /workspace
 
+COPY tools/setup_deepgemm_pythons.sh tools/build_deepgemm_C.py tools/
+ENV DEEPGEMM_VENV_PREFIX=/opt/dgenv
+RUN --mount=type=cache,target=/root/.cache/uv \
+    tools/setup_deepgemm_pythons.sh > /tmp/deepgemm-pythons.txt
+
 COPY . .
 
+# setup.py packages these separately built Rust artifacts as-is.
+RUN rm -f aphrodite/aphrodite-rs aphrodite/_rust_*.so
+COPY --from=rust-build /workspace/aphrodite/aphrodite-rs aphrodite/aphrodite-rs
+COPY --from=rust-build /workspace/aphrodite/_rust_*.so aphrodite/
+
 ARG GIT_REPO_CHECK=0
 RUN --mount=type=bind,source=.git,target=.git \
     if [ "$GIT_REPO_CHECK" != "0" ]; then bash tools/check_repo.sh ; fi
@@ -318,6 +394,7 @@ ENV APHRODITE_VERSION_OVERRIDE=${APHRODITE_VERSION_OVERRIDE}
 ENV MAX_JOBS=${max_jobs}
 ENV NVCC_THREADS=${nvcc_threads}
 ENV CMAKE_BUILD_TYPE=Release
+ENV APHRODITE_USE_PRECOMPILED_RUST=1
 
 # Use existing torch for nightly builds
 RUN --mount=type=cache,target=/root/.cache/uv \
@@ -336,11 +413,15 @@ RUN --mount=type=cache,id=aphrodite-wheel-uv,target=/root/.cache/uv \
         export RUSTC_WRAPPER=sccache; \
         sccache --show-stats; \
     fi \
+    && export DEEPGEMM_PYTHON_INTERPRETERS="$(cat /tmp/deepgemm-pythons.txt)" \
     && python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 \
     && if [ "$USE_SCCACHE" = "1" ]; then \
         sccache --show-stats; \
     fi
 
+RUN python3 .github/scripts/verify_cuda_wheel_artifacts.py \
+    "$(find dist -maxdepth 1 -type f -name '*.whl' -print -quit)"
+
 # Check the size of the CUDA 13 wheel uploaded to PyPI
 COPY .buildkite/check-wheel-size.py check-wheel-size.py
 # sync the default value with .buildkite/check-wheel-size.py

docker/Dockerfile.rocm

@@ -193,7 +193,7 @@ RUN --mount=type=cache,id=aphrodite-rocm-cargo-registry,target=/root/.cargo/regi
         export RUSTC_WRAPPER=sccache; \
         sccache --show-stats; \
     fi \
-    && python3 tools/build_rust.py \
+    && python3 tools/build_rust.py --release \
     && if [ "$USE_SCCACHE" = "1" ]; then \
         sccache --show-stats; \
     fi

tools/setup_deepgemm_pythons.sh

@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Provision one bare Python environment per supported CPython version and
+# print their interpreter paths as a colon-separated list for CMake.
+#
+# Usage:
+#   export DEEPGEMM_PYTHON_INTERPRETERS="$(
+#     tools/setup_deepgemm_pythons.sh
+#   )"
+#
+# Optional: DEEPGEMM_VENV_PREFIX (default: /tmp/dgenv).
+set -euo pipefail
+
+if [ "$#" -eq 0 ]; then
+  pyproject="$(dirname "$0")/../pyproject.toml"
+  spec="$(
+    grep -E '^requires-python' "$pyproject" |
+      grep -oE '>=3\.[0-9]+,<3\.[0-9]+'
+  )"
+  lo="${spec#>=3.}"
+  lo="${lo%%,*}"
+  hi="${spec##*<3.}"
+  set -- $(seq "$lo" $((hi - 1)) | sed 's/^/3./')
+fi
+
+prefix="${DEEPGEMM_VENV_PREFIX:-/tmp/dgenv}"
+mkdir -p "$prefix"
+
+paths=""
+for version in "$@"; do
+  venv="${prefix}/${version}"
+  if [ ! -x "${venv}/bin/python" ]; then
+    uv venv \
+      --python "$version" \
+      "$venv" \
+      --python-preference only-managed \
+      --seed \
+      >/dev/null
+  fi
+  paths="${paths}:${venv}/bin/python"
+done
+
+echo "${paths#:}"