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, 19:00 UTC
[Build] DeepGEMM: trim comments, add integration notes + TODOs (#42429)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
May 12, 2026 at 18:57 UTC
184577ae46f5682c28725eb502ef5448d09080da
4 files changed
+40
-43
cmake/external_projects/deepgemm.cmake
+27
-14
@@ -53,14 +53,30 @@ cuda_archs_loose_intersection(DEEPGEMM_ARCHS
53
if(DEEPGEMM_ARCHS)
54
message(STATUS "DeepGEMM CUDA architectures: ${DEEPGEMM_ARCHS}")
55
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.
56
+ #
57
+ # DeepGEMM integration notes
58
+ # --------------------------
59
+ # We vendor DeepGEMM into vllm/third_party/deep_gemm/ and bundle a
60
+ # `_C.cpython-X.Y-*.so` for every CPython in `requires-python`. The
61
+ # per-Python build is delegated to tools/build_deepgemm_C.py.
62
+ #
63
+ # Why per-Python: DeepGEMM's binding uses PYBIND11_MODULE, which links
64
+ # private CPython symbols — a single `_C.abi3.so` is not viable today
65
+ # (see #41476 / #41512 for the failed attempt).
66
+ #
67
+ # TODOs (tracked in vllm-project/vllm#42431):
68
+ # - Replace DeepGEMM's pybind11 binding with a TORCH_LIBRARY + shim
69
+ # binding (cf. vllm-flash-attention/csrc/common/pytorch_shim.h) to
70
+ # collapse to one `_C.abi3.so`. Needs either an upstream change or
71
+ # a maintained binding fork in vLLM.
72
+ # - AOT-compile DeepGEMM's CUDA kernels instead of runtime JIT to drop
73
+ # the vendored CUTLASS/CCCL headers and the CUDA-toolkit-at-runtime
74
+ # requirement.
75
+ #
76
+
77
+ # DEEPGEMM_PYTHON_INTERPRETERS: ":"-separated target Python paths.
78
+ # Empty/unset → fall back to the build interpreter (editable installs).
79
+ # (Empty-but-set env vars test as DEFINED in cmake — treat as unset.)
80
if(NOT "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}" STREQUAL "")
81
string(REPLACE ":" ";" _dg_pythons "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}")
82
else()
@@ -68,10 +84,8 @@ if(DEEPGEMM_ARCHS)
84
endif()
85
message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}")
86
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.
87
+ # add_custom_command does no implicit header scanning; glob explicitly so
88
+ # header-only edits in DeepGEMM/cutlass/fmt re-trigger the rebuild.
89
file(GLOB_RECURSE _dg_headers
90
"${deepgemm_SOURCE_DIR}/csrc/*.h"
91
"${deepgemm_SOURCE_DIR}/csrc/*.hpp"
@@ -88,8 +102,7 @@ if(DEEPGEMM_ARCHS)
102
OUTPUT_VARIABLE _dg_soabi
103
OUTPUT_STRIP_TRAILING_WHITESPACE
104
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.
105
+ # Dedup interpreters that resolve to the same CPython.
106
if(_dg_soabi IN_LIST _dg_seen_soabis)
107
continue()
108
endif()
docker/Dockerfile
+3
-4
@@ -301,10 +301,9 @@ 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.
304
+# Provision one bare Python per `requires-python` entry; cmake reads
305
+# DEEPGEMM_PYTHON_INTERPRETERS to build DeepGEMM `_C` for each. See
306
+# cmake/external_projects/deepgemm.cmake for the full picture.
307
COPY tools/setup_deepgemm_pythons.sh tools/build_deepgemm_C.py tools/
308
ENV DEEPGEMM_VENV_PREFIX=/opt/dgenv
309
RUN --mount=type=cache,target=/root/.cache/uv \
tools/build_deepgemm_C.py
+5
-7
@@ -1,11 +1,10 @@
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.
3
+"""Build DeepGEMM's `_C` pybind11 extension for <TARGET_PY>.
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.
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
"""
@@ -42,8 +41,7 @@ info = json.loads(
41
cuda_home = cpp_extension.CUDA_HOME
42
if cuda_home is None:
43
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.
44
+# CCCL lives outside the standard CUDAToolkit search (mirrors DeepGEMM's setup.py).
45
includes = [
46
info["INCLUDEPY"],
47
f"{cuda_home}/include",
tools/setup_deepgemm_pythons.sh
+5
-18
@@ -1,23 +1,12 @@
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.
2
+# Provision one bare Python per `requires-python` entry (or per argument) and
3
+# print their paths as ":"-separated DEEPGEMM_PYTHON_INTERPRETERS. Skip this
4
+# entirely if you already have interpreter paths.
5
#
6
# Usage:
7
# export DEEPGEMM_PYTHON_INTERPRETERS=$(tools/setup_deepgemm_pythons.sh)
8
# python setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
9
#
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
-#
10
# Optional: DEEPGEMM_VENV_PREFIX (default: /tmp/dgenv).
11
set -euo pipefail
12
@@ -37,10 +26,8 @@ mkdir -p "$prefix"
26
paths=""
27
for V in "$@"; do
28
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.
29
+ # uv-managed Python ensures Python.h is present; system 3.X-dev packages
30
+ # on the manylinux / Ubuntu build bases are not always installed.
31
[ -x "$venv/bin/python" ] || \
32
uv venv --python "$V" "$venv" --python-preference only-managed --seed \
33
>/dev/null