[Build] Build bundled DeepGEMM _C per-Python so the wheel imports on every CPython
#41516
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,87 @@ | ||
| # SPDX-License-Identifier: Apache-2.0 | ||
| # SPDX-FileCopyrightText: Copyright contributors to the vLLM project | ||
| """Build DeepGEMM's `_C` pybind11 extension for a target Python. | ||
|
|
||
| Driven from `cmake/external_projects/deepgemm.cmake`. The driver is the | ||
| build interpreter (which has torch); the *target* Python is only used for | ||
| its header path and SOABI. This avoids needing torch installed in N venvs | ||
| to produce N matching `.so` files. | ||
|
|
||
| Usage: python build_deepgemm_C.py <DEEPGEMM_SRC_DIR> <OUTPUT_DIR> <TARGET_PY> | ||
| """ | ||
|
|
||
| import json | ||
| import os | ||
| import subprocess | ||
| import sys | ||
|
mgoin marked this conversation as resolved.
|
||
| from pathlib import Path | ||
|
|
||
| import torch | ||
| from torch.utils import cpp_extension | ||
|
|
||
| if len(sys.argv) != 4: | ||
| sys.exit(f"usage: {sys.argv[0]} <SRC> <OUT> <TARGET_PY>") | ||
|
|
||
| src = Path(sys.argv[1]).resolve() | ||
| out = Path(sys.argv[2]).resolve() | ||
| target_py = sys.argv[3] | ||
| out.mkdir(parents=True, exist_ok=True) | ||
|
|
||
| info = json.loads( | ||
| subprocess.check_output( | ||
| [ | ||
| target_py, | ||
| "-c", | ||
| "import sysconfig, json; " | ||
| "print(json.dumps({k: sysconfig.get_config_var(k) " | ||
| "for k in ('EXT_SUFFIX', 'INCLUDEPY')}))", | ||
| ] | ||
| ).decode() | ||
| ) | ||
|
|
||
| cuda_home = cpp_extension.CUDA_HOME | ||
|
mgoin marked this conversation as resolved.
|
||
| if cuda_home is None: | ||
| sys.exit("CUDA_HOME not found; cannot build DeepGEMM _C") | ||
| # CCCL lives outside the standard CUDAToolkit search, mirroring DeepGEMM's | ||
| # own setup.py. | ||
| includes = [ | ||
| info["INCLUDEPY"], | ||
| f"{cuda_home}/include", | ||
| f"{cuda_home}/include/cccl", | ||
| str(src / "csrc"), | ||
| str(src / "deep_gemm/include"), | ||
| str(src / "third-party/cutlass/include"), | ||
| str(src / "third-party/cutlass/tools/util/include"), | ||
| str(src / "third-party/fmt/include"), | ||
| *cpp_extension.include_paths(device_type="cuda"), | ||
| ] | ||
|
|
||
| cmd = [ | ||
| os.environ.get("CXX", "g++"), | ||
| "-shared", | ||
| "-fPIC", | ||
| "-std=c++17", | ||
| "-O3", | ||
| "-g0", | ||
| "-Wno-psabi", | ||
| "-Wno-deprecated-declarations", | ||
| "-DTORCH_API_INCLUDE_EXTENSION_H", | ||
| "-DTORCH_EXTENSION_NAME=_C", | ||
| f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}", | ||
| *(f"-I{p}" for p in includes), | ||
| str(src / "csrc/python_api.cpp"), | ||
| *(f"-L{p}" for p in cpp_extension.library_paths(device_type="cuda")), | ||
| f"-L{cuda_home}/lib64", | ||
| "-ltorch", | ||
| "-ltorch_python", | ||
| "-ltorch_cpu", | ||
| "-ltorch_cuda", | ||
| "-lc10", | ||
| "-lc10_cuda", | ||
| "-lcudart", | ||
| "-lnvrtc", | ||
| "-o", | ||
| str(out / f"_C{info['EXT_SUFFIX']}"), | ||
| ] | ||
| print("[build_deepgemm_C] " + " ".join(cmd), flush=True) | ||
| subprocess.check_call(cmd) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,38 @@ | ||
| # SPDX-License-Identifier: Apache-2.0 | ||
| # SPDX-FileCopyrightText: Copyright contributors to the vLLM project | ||
|
|
||
| """Assert the installed vLLM has a `_C.cpython-X.Y-*.so` for every CPython | ||
| covered by `requires-python`. Fails closed if a Python's `.so` is missing | ||
| from the wheel — i.e. the regression that surfaced in #41476/#41512. | ||
|
|
||
| Run from a CI test job after vLLM is installed, e.g. the H100 deepgemm | ||
| kernel tests in .buildkite/test_areas/kernels.yaml. | ||
| """ | ||
|
|
||
| import os | ||
| import sys | ||
| from pathlib import Path | ||
|
|
||
| import regex as re | ||
|
mgoin marked this conversation as resolved.
|
||
| import tomllib | ||
|
|
||
| import vllm.third_party.deep_gemm as pkg | ||
|
Check warning on line 19 in tools/check_wheel_deepgemm.py
|
||
|
mgoin marked this conversation as resolved.
Outdated
|
||
|
|
||
| SO_RE = re.compile(r"^_C\.cpython-(\d)(\d+)-") | ||
|
|
||
|
|
||
| def required_pythons() -> list[str]: | ||
| pyproject = Path(__file__).resolve().parent.parent / "pyproject.toml" | ||
| spec = tomllib.loads(pyproject.read_text())["project"]["requires-python"] | ||
| m = re.match(r">=3\.(\d+),<3\.(\d+)", spec) | ||
| if not m: | ||
| sys.exit(f"unexpected requires-python format: {spec!r}") | ||
| return [f"3.{v}" for v in range(int(m[1]), int(m[2]))] | ||
|
|
||
|
|
||
| pkg_dir = Path(pkg.__file__).parent | ||
| found = {f"{m[1]}.{m[2]}" for f in os.listdir(pkg_dir) if (m := SO_RE.match(f))} | ||
| required = required_pythons() | ||
| missing = [v for v in required if v not in found] | ||
| print(f"deepgemm _C: found {sorted(found)}, required {required}, missing {missing}") | ||
| sys.exit(1 if missing else 0) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,43 @@ | ||
| #!/usr/bin/env bash | ||
| # Provision bare Python interpreters for the DeepGEMM `_C` per-Python build | ||
| # and print a colon-separated list of their paths to stdout. | ||
| # | ||
| # Each target Python only needs a working interpreter — torch is not | ||
| # installed since `tools/build_deepgemm_C.py` runs from the build interpreter. | ||
| # uv re-uses any matching system Python and downloads a managed build | ||
| # otherwise. | ||
| # | ||
| # Usage: | ||
| # export DEEPGEMM_PYTHON_INTERPRETERS=$(tools/setup_deepgemm_pythons.sh) | ||
| # python setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 | ||
| # | ||
| # With no args, expands to every CPython covered by `requires-python` in | ||
| # pyproject.toml. Pass explicit versions (e.g. `3.10 3.11`) to override. | ||
| # | ||
| # Skip this script if you don't have uv: set DEEPGEMM_PYTHON_INTERPRETERS | ||
| # directly to existing interpreter paths. Editable / single-Python builds | ||
| # don't need the env var at all (cmake falls back to the build interpreter). | ||
| # | ||
| # Optional: DEEPGEMM_VENV_PREFIX (default: /tmp/dgenv). | ||
| set -euo pipefail | ||
|
|
||
| if [ "$#" -eq 0 ]; then | ||
| # Derive the matrix from `requires-python = ">=3.X,<3.Y"` in pyproject.toml. | ||
| 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 V in "$@"; do | ||
| venv="$prefix/$V" | ||
| [ -x "$venv/bin/python" ] || uv venv --python "$V" "$venv" --seed >/dev/null | ||
| paths="$paths:$venv/bin/python" | ||
| done | ||
| echo "${paths#:}" |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.