| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Build DeepGEMM's `_C` pybind11 extension for <TARGET_PY>. |
| 4 | |
| 5 | Driven from cmake/external_projects/deepgemm.cmake. The driver runs against |
| 6 | the build interpreter's torch; <TARGET_PY> is only consulted for INCLUDEPY |
| 7 | and SOABI, so target venvs don't need torch installed. |
| 8 | |
| 9 | Usage: python build_deepgemm_C.py <DEEPGEMM_SRC_DIR> <OUTPUT_DIR> <TARGET_PY> |
| 10 | """ |
| 11 | |
| 12 | import json |
| 13 | import os |
| 14 | import subprocess |
| 15 | import sys |
| 16 | from pathlib import Path |
| 17 | |
| 18 | import torch |
| 19 | from torch.utils import cpp_extension |
| 20 | |
| 21 | if len(sys.argv) != 4: |
| 22 | sys.exit(f"usage: {sys.argv[0]} <SRC> <OUT> <TARGET_PY>") |
| 23 | |
| 24 | src = Path(sys.argv[1]).resolve() |
| 25 | out = Path(sys.argv[2]).resolve() |
| 26 | target_py = sys.argv[3] |
| 27 | out.mkdir(parents=True, exist_ok=True) |
| 28 | |
| 29 | info = json.loads( |
| 30 | subprocess.check_output( |
| 31 | [ |
| 32 | target_py, |
| 33 | "-c", |
| 34 | "import sysconfig, json; " |
| 35 | "print(json.dumps({k: sysconfig.get_config_var(k) " |
| 36 | "for k in ('EXT_SUFFIX', 'INCLUDEPY')}))", |
| 37 | ] |
| 38 | ).decode() |
| 39 | ) |
| 40 | |
| 41 | cuda_home = cpp_extension.CUDA_HOME |
| 42 | if cuda_home is None: |
| 43 | sys.exit("CUDA_HOME not found; cannot build DeepGEMM _C") |
| 44 | # CCCL lives outside the standard CUDAToolkit search (mirrors DeepGEMM's setup.py). |
| 45 | includes = [ |
| 46 | info["INCLUDEPY"], |
| 47 | f"{cuda_home}/include", |
| 48 | f"{cuda_home}/include/cccl", |
| 49 | str(src / "csrc"), |
| 50 | str(src / "deep_gemm/include"), |
| 51 | str(src / "third-party/cutlass/include"), |
| 52 | str(src / "third-party/cutlass/tools/util/include"), |
| 53 | str(src / "third-party/deep_jit/include"), |
| 54 | *cpp_extension.include_paths(device_type="cuda"), |
| 55 | ] |
| 56 | |
| 57 | cmd = [ |
| 58 | os.environ.get("CXX", "g++"), |
| 59 | "-shared", |
| 60 | "-fPIC", |
| 61 | "-std=c++20", |
| 62 | "-O3", |
| 63 | "-g0", |
| 64 | "-Wno-psabi", |
| 65 | "-Wno-deprecated-declarations", |
| 66 | "-DTORCH_API_INCLUDE_EXTENSION_H", |
| 67 | "-DTORCH_EXTENSION_NAME=_C", |
| 68 | f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}", |
| 69 | *(f"-I{p}" for p in includes), |
| 70 | str(src / "csrc/python_api.cpp"), |
| 71 | *(f"-L{p}" for p in cpp_extension.library_paths(device_type="cuda")), |
| 72 | f"-L{cuda_home}/lib64", |
| 73 | "-ltorch", |
| 74 | "-ltorch_python", |
| 75 | "-ltorch_cpu", |
| 76 | "-ltorch_cuda", |
| 77 | "-lc10", |
| 78 | "-lc10_cuda", |
| 79 | "-lcudart", |
| 80 | "-lnvrtc", |
| 81 | "-o", |
| 82 | str(out / f"_C{info['EXT_SUFFIX']}"), |
| 83 | ] |
| 84 | print("[build_deepgemm_C] " + " ".join(cmd), flush=True) |
| 85 | subprocess.check_call(cmd) |