A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 20:15 UTC
[Build] Build bundled DeepGEMM `_C` per-Python so the wheel imports on every CPython (#41516)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Michael Goin committed
May 12, 2026 at 10:27 UTC
d077622d600322d46abe492ea8355a871ba3a2fc
6 files changed
+254
-42
.buildkite/test_areas/kernels.yaml
+5
@@ -105,7 +105,11 @@ steps:
105
device: h100
106
num_devices: 1
107
source_file_dependencies:
108
+ - cmake/external_projects/deepgemm.cmake
109
- tools/install_deepgemm.sh
110
+ - tools/build_deepgemm_C.py
111
+ - tools/setup_deepgemm_pythons.sh
112
+ - tools/check_wheel_deepgemm.py
113
- vllm/utils/deep_gemm.py
114
- vllm/model_executor/layers/fused_moe
115
- vllm/model_executor/layers/quantization
@@ -115,6 +119,7 @@ steps:
119
- tests/kernels/attention/test_deepgemm_attention.py
120
- tests/quantization/test_cutlass_w4a16.py
121
commands:
122
+ - python3 ../tools/check_wheel_deepgemm.py
123
- pytest -v -s kernels/quantization/test_block_fp8.py
124
- pytest -v -s kernels/moe/test_deepgemm.py
125
- pytest -v -s kernels/moe/test_batched_deepgemm.py
cmake/external_projects/deepgemm.cmake
+61
-42
@@ -53,48 +53,67 @@ cuda_archs_loose_intersection(DEEPGEMM_ARCHS
53
if(DEEPGEMM_ARCHS)
54
message(STATUS "DeepGEMM CUDA architectures: ${DEEPGEMM_ARCHS}")
55
56
- find_package(CUDAToolkit REQUIRED)
57
-
58
- #
59
- # Build the _C pybind11 extension from DeepGEMM's C++ source.
60
- # This is a CXX-only module — CUDA kernels are JIT-compiled at runtime.
61
- #
62
- Python_add_library(_deep_gemm_C MODULE WITH_SOABI
63
- "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp")
64
-
65
- # The pybind11 module name must be _C to match DeepGEMM's Python imports.
66
- set_target_properties(_deep_gemm_C PROPERTIES OUTPUT_NAME "_C")
67
-
68
- target_compile_definitions(_deep_gemm_C PRIVATE
69
- "-DTORCH_EXTENSION_NAME=_C")
70
-
71
- target_include_directories(_deep_gemm_C PRIVATE
72
- "${deepgemm_SOURCE_DIR}/csrc"
73
- "${deepgemm_SOURCE_DIR}/deep_gemm/include"
74
- "${deepgemm_SOURCE_DIR}/third-party/cutlass/include"
75
- "${deepgemm_SOURCE_DIR}/third-party/cutlass/tools/util/include"
76
- "${deepgemm_SOURCE_DIR}/third-party/fmt/include")
77
-
78
- target_compile_options(_deep_gemm_C PRIVATE
79
- $<$<COMPILE_LANGUAGE:CXX>:-O3>
80
- $<$<COMPILE_LANGUAGE:CXX>:-Wno-psabi>
81
- $<$<COMPILE_LANGUAGE:CXX>:-Wno-deprecated-declarations>)
82
-
83
- # torch_python is required because DeepGEMM uses pybind11 type casters
84
- # for at::Tensor (via PYBIND11_MODULE), unlike vLLM's own extensions which
85
- # use torch::Library custom ops.
86
- find_library(TORCH_PYTHON_LIBRARY torch_python
87
- PATHS "${TORCH_INSTALL_PREFIX}/lib"
88
- REQUIRED)
89
-
90
- target_link_libraries(_deep_gemm_C PRIVATE
91
- torch ${TORCH_LIBRARIES} "${TORCH_PYTHON_LIBRARY}"
92
- CUDA::cudart CUDA::nvrtc)
93
-
94
- # Install the shared library into the vendored package directory
95
- install(TARGETS _deep_gemm_C
96
- LIBRARY DESTINATION vllm/third_party/deep_gemm
97
- COMPONENT _deep_gemm_C)
56
+ # Build _C once per interpreter in DEEPGEMM_PYTHON_INTERPRETERS (":"-
57
+ # separated paths) so the wheel imports cleanly on every supported Python.
58
+ # Unset → fall back to the build interpreter (editable / source builds).
59
+ # The compile is delegated to tools/build_deepgemm_C.py and always runs
60
+ # against the build interpreter's torch — target Pythons don't need torch.
61
+ # Note: empty-but-set env vars are still DEFINED in cmake; treat empty as
62
+ # unset so an empty interpreter list falls back to the build interpreter
63
+ # rather than silently skipping the per-Python build.
64
+ if(NOT "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}" STREQUAL "")
65
+ string(REPLACE ":" ";" _dg_pythons "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}")
66
+ else()
67
+ set(_dg_pythons "${Python_EXECUTABLE}")
68
+ endif()
69
+ message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}")
70
+
71
+ # Header set fed to add_custom_command's DEPENDS so a header-only edit
72
+ # (in upstream DeepGEMM or its vendored cutlass/fmt) re-triggers the
73
+ # rebuild. add_custom_command does no implicit header scanning, unlike
74
+ # add_library.
75
+ file(GLOB_RECURSE _dg_headers
76
+ "${deepgemm_SOURCE_DIR}/csrc/*.h"
77
+ "${deepgemm_SOURCE_DIR}/csrc/*.hpp"
78
+ "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.h"
79
+ "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.hpp"
80
+ "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh")
81
+
82
+ set(_dg_markers)
83
+ set(_dg_seen_soabis)
84
+ foreach(_pybin IN LISTS _dg_pythons)
85
+ execute_process(
86
+ COMMAND "${_pybin}" -c
87
+ "import sysconfig; print(sysconfig.get_config_var('SOABI'))"
88
+ OUTPUT_VARIABLE _dg_soabi
89
+ OUTPUT_STRIP_TRAILING_WHITESPACE
90
+ COMMAND_ERROR_IS_FATAL ANY)
91
+ # Dedup so duplicate paths (or two paths resolving to the same CPython)
92
+ # don't register conflicting build rules.
93
+ if(_dg_soabi IN_LIST _dg_seen_soabis)
94
+ continue()
95
+ endif()
96
+ list(APPEND _dg_seen_soabis "${_dg_soabi}")
97
+ set(_dg_dir "${CMAKE_CURRENT_BINARY_DIR}/deepgemm_C_${_dg_soabi}")
98
+ set(_dg_marker "${_dg_dir}/.built")
99
+ add_custom_command(
100
+ OUTPUT "${_dg_marker}"
101
+ COMMAND "${Python_EXECUTABLE}"
102
+ "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py"
103
+ "${deepgemm_SOURCE_DIR}" "${_dg_dir}" "${_pybin}"
104
+ COMMAND "${CMAKE_COMMAND}" -E touch "${_dg_marker}"
105
+ DEPENDS "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py"
106
+ "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp"
107
+ ${_dg_headers}
108
+ COMMENT "Building DeepGEMM _C for ${_pybin}"
109
+ VERBATIM)
110
+ list(APPEND _dg_markers "${_dg_marker}")
111
+ install(DIRECTORY "${_dg_dir}/"
112
+ DESTINATION vllm/third_party/deep_gemm
113
+ COMPONENT _deep_gemm_C
114
+ FILES_MATCHING PATTERN "_C.cpython-*.so")
115
+ endforeach()
116
+ add_custom_target(_deep_gemm_C ALL DEPENDS ${_dg_markers})
117
118
#
119
# Vendor DeepGEMM Python package files
docker/Dockerfile
+11
@@ -301,6 +301,15 @@ RUN --mount=type=cache,target=/root/.cache/uv \
301
python3 use_existing_torch.py --prefix; \
302
fi
303
304
+# Provision a bare interpreter for each CPython covered by `requires-python`
305
+# so DeepGEMM `_C` is built once per Python and bundled side-by-side in the
306
+# wheel; cmake reads DEEPGEMM_PYTHON_INTERPRETERS in deepgemm.cmake's
307
+# foreach loop. The matrix is derived from pyproject.toml.
308
+COPY tools/setup_deepgemm_pythons.sh tools/build_deepgemm_C.py tools/
309
+ENV DEEPGEMM_VENV_PREFIX=/opt/dgenv
310
+RUN --mount=type=cache,target=/root/.cache/uv \
311
+ tools/setup_deepgemm_pythons.sh > /tmp/dg_pythons.txt
312
+
313
# Build the vLLM wheel
314
# if USE_SCCACHE is set, use sccache to speed up compilation
315
# AWS credentials mounted at ~/.aws/credentials for sccache S3 auth (optional)
@@ -328,6 +337,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
337
&& export VLLM_PRECOMPILED_WHEEL_COMMIT="${VLLM_MERGE_BASE_COMMIT}" \
338
&& export VLLM_MAIN_CUDA_VERSION="${VLLM_MAIN_CUDA_VERSION}" \
339
&& export VLLM_DOCKER_BUILD_CONTEXT=1 \
340
+ && export DEEPGEMM_PYTHON_INTERPRETERS=$(cat /tmp/dg_pythons.txt) \
341
&& sccache --show-stats \
342
&& python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 \
343
&& sccache --show-stats; \
@@ -345,6 +355,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
355
export VLLM_USE_PRECOMPILED="${VLLM_USE_PRECOMPILED}" && \
356
export VLLM_PRECOMPILED_WHEEL_COMMIT="${VLLM_MERGE_BASE_COMMIT}" && \
357
export VLLM_DOCKER_BUILD_CONTEXT=1 && \
358
+ export DEEPGEMM_PYTHON_INTERPRETERS=$(cat /tmp/dg_pythons.txt) && \
359
python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
360
fi
361
tools/build_deepgemm_C.py
new
+87
@@ -0,0 +1,87 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Build DeepGEMM's `_C` pybind11 extension for a target Python.
4
+
5
+Driven from `cmake/external_projects/deepgemm.cmake`. The driver is the
6
+build interpreter (which has torch); the *target* Python is only used for
7
+its header path and SOABI. This avoids needing torch installed in N venvs
8
+to produce N matching `.so` files.
9
+
10
+Usage: python build_deepgemm_C.py <DEEPGEMM_SRC_DIR> <OUTPUT_DIR> <TARGET_PY>
11
+"""
12
+
13
+import json
14
+import os
15
+import subprocess
16
+import sys
17
+from pathlib import Path
18
+
19
+import torch
20
+from torch.utils import cpp_extension
21
+
22
+if len(sys.argv) != 4:
23
+ sys.exit(f"usage: {sys.argv[0]} <SRC> <OUT> <TARGET_PY>")
24
+
25
+src = Path(sys.argv[1]).resolve()
26
+out = Path(sys.argv[2]).resolve()
27
+target_py = sys.argv[3]
28
+out.mkdir(parents=True, exist_ok=True)
29
+
30
+info = json.loads(
31
+ subprocess.check_output(
32
+ [
33
+ target_py,
34
+ "-c",
35
+ "import sysconfig, json; "
36
+ "print(json.dumps({k: sysconfig.get_config_var(k) "
37
+ "for k in ('EXT_SUFFIX', 'INCLUDEPY')}))",
38
+ ]
39
+ ).decode()
40
+)
41
+
42
+cuda_home = cpp_extension.CUDA_HOME
43
+if cuda_home is None:
44
+ sys.exit("CUDA_HOME not found; cannot build DeepGEMM _C")
45
+# CCCL lives outside the standard CUDAToolkit search, mirroring DeepGEMM's
46
+# own setup.py.
47
+includes = [
48
+ info["INCLUDEPY"],
49
+ f"{cuda_home}/include",
50
+ f"{cuda_home}/include/cccl",
51
+ str(src / "csrc"),
52
+ str(src / "deep_gemm/include"),
53
+ str(src / "third-party/cutlass/include"),
54
+ str(src / "third-party/cutlass/tools/util/include"),
55
+ str(src / "third-party/fmt/include"),
56
+ *cpp_extension.include_paths(device_type="cuda"),
57
+]
58
+
59
+cmd = [
60
+ os.environ.get("CXX", "g++"),
61
+ "-shared",
62
+ "-fPIC",
63
+ "-std=c++20",
64
+ "-O3",
65
+ "-g0",
66
+ "-Wno-psabi",
67
+ "-Wno-deprecated-declarations",
68
+ "-DTORCH_API_INCLUDE_EXTENSION_H",
69
+ "-DTORCH_EXTENSION_NAME=_C",
70
+ f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}",
71
+ *(f"-I{p}" for p in includes),
72
+ str(src / "csrc/python_api.cpp"),
73
+ *(f"-L{p}" for p in cpp_extension.library_paths(device_type="cuda")),
74
+ f"-L{cuda_home}/lib64",
75
+ "-ltorch",
76
+ "-ltorch_python",
77
+ "-ltorch_cpu",
78
+ "-ltorch_cuda",
79
+ "-lc10",
80
+ "-lc10_cuda",
81
+ "-lcudart",
82
+ "-lnvrtc",
83
+ "-o",
84
+ str(out / f"_C{info['EXT_SUFFIX']}"),
85
+]
86
+print("[build_deepgemm_C] " + " ".join(cmd), flush=True)
87
+subprocess.check_call(cmd)
tools/check_wheel_deepgemm.py
new
+41
@@ -0,0 +1,41 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Assert the installed vLLM has a `_C.cpython-X.Y-*.so` for every CPython
5
+covered by `requires-python`. Fails closed if a Python's `.so` is missing
6
+from the wheel — i.e. the regression that surfaced in #41476/#41512.
7
+
8
+Run from a CI test job after vLLM is installed, e.g. the H100 deepgemm
9
+kernel tests in .buildkite/test_areas/kernels.yaml.
10
+"""
11
+
12
+import importlib.util
13
+import os
14
+import sys
15
+from pathlib import Path
16
+
17
+import regex as re
18
+import tomllib
19
+
20
+SO_RE = re.compile(r"^_C\.cpython-(\d)(\d+)-")
21
+
22
+
23
+def required_pythons() -> list[str]:
24
+ pyproject = Path(__file__).resolve().parent.parent / "pyproject.toml"
25
+ spec = tomllib.loads(pyproject.read_text())["project"]["requires-python"]
26
+ m = re.match(r">=3\.(\d+),<3\.(\d+)", spec)
27
+ if not m:
28
+ sys.exit(f"unexpected requires-python format: {spec!r}")
29
+ return [f"3.{v}" for v in range(int(m[1]), int(m[2]))]
30
+
31
+
32
+spec = importlib.util.find_spec("vllm.third_party.deep_gemm")
33
+if spec is None or spec.origin is None:
34
+ sys.exit("vllm.third_party.deep_gemm not importable; is vllm installed?")
35
+pkg_dir = Path(spec.origin).parent
36
+
37
+found = {f"{m[1]}.{m[2]}" for f in os.listdir(pkg_dir) if (m := SO_RE.match(f))}
38
+required = required_pythons()
39
+missing = [v for v in required if v not in found]
40
+print(f"deepgemm _C: found {sorted(found)}, required {required}, missing {missing}")
41
+sys.exit(1 if missing else 0)
tools/setup_deepgemm_pythons.sh
new
+49
@@ -0,0 +1,49 @@
1
+#!/usr/bin/env bash
2
+# Provision bare Python interpreters for the DeepGEMM `_C` per-Python build
3
+# and print a colon-separated list of their paths to stdout.
4
+#
5
+# Each target Python only needs a working interpreter — torch is not
6
+# installed since `tools/build_deepgemm_C.py` runs from the build interpreter.
7
+# uv re-uses any matching system Python and downloads a managed build
8
+# otherwise.
9
+#
10
+# Usage:
11
+# export DEEPGEMM_PYTHON_INTERPRETERS=$(tools/setup_deepgemm_pythons.sh)
12
+# python setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
13
+#
14
+# With no args, expands to every CPython covered by `requires-python` in
15
+# pyproject.toml. Pass explicit versions (e.g. `3.10 3.11`) to override.
16
+#
17
+# Skip this script if you don't have uv: set DEEPGEMM_PYTHON_INTERPRETERS
18
+# directly to existing interpreter paths. Editable / single-Python builds
19
+# don't need the env var at all (cmake falls back to the build interpreter).
20
+#
21
+# Optional: DEEPGEMM_VENV_PREFIX (default: /tmp/dgenv).
22
+set -euo pipefail
23
+
24
+if [ "$#" -eq 0 ]; then
25
+ # Derive the matrix from `requires-python = ">=3.X,<3.Y"` in pyproject.toml.
26
+ pyproject="$(dirname "$0")/../pyproject.toml"
27
+ spec=$(grep -E '^requires-python' "$pyproject" \
28
+ | grep -oE '>=3\.[0-9]+,<3\.[0-9]+')
29
+ lo=${spec#>=3.}; lo=${lo%%,*}
30
+ hi=${spec##*<3.}
31
+ set -- $(seq "$lo" $((hi - 1)) | sed 's/^/3./')
32
+fi
33
+
34
+prefix="${DEEPGEMM_VENV_PREFIX:-/tmp/dgenv}"
35
+mkdir -p "$prefix"
36
+
37
+paths=""
38
+for V in "$@"; do
39
+ venv="$prefix/$V"
40
+ # Force a managed (uv-downloaded) Python so dev headers are bundled.
41
+ # System Pythons on the build base may lack headers (manylinux's
42
+ # /opt/python/cpXY-cpXY are off PATH; an apt-installed python3.X often
43
+ # has no -dev), and the per-Python build needs Python.h.
44
+ [ -x "$venv/bin/python" ] || \
45
+ uv venv --python "$V" "$venv" --python-preference only-managed --seed \
46
+ >/dev/null
47
+ paths="$paths:$venv/bin/python"
48
+done
49
+echo "${paths#:}"