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 5, 09:45 UTC
[UX] Integrate DeepGEMM into vLLM wheel via CMake (#37980)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Michael Goin committed
Apr 9, 2026 at 03:56 UTC
eb4205fee52dce68f916e1793faec1fa557a014d
12 files changed
+251
-40
.gitignore
+3
@@ -12,6 +12,9 @@ vllm/third_party/triton_kernels/*
12
# FlashMLA interface copied from source
13
vllm/third_party/flashmla/flash_mla_interface.py
14
15
+# DeepGEMM vendored package built from source
16
+vllm/third_party/deep_gemm/
17
+
18
# triton jit
19
.triton
20
CMakeLists.txt
+1
@@ -1222,6 +1222,7 @@ endif()
1222
1223
# For CUDA we also build and ship some external projects.
1224
if (VLLM_GPU_LANG STREQUAL "CUDA")
1225
+ include(cmake/external_projects/deepgemm.cmake)
1226
include(cmake/external_projects/flashmla.cmake)
1227
include(cmake/external_projects/qutlass.cmake)
1228
cmake/external_projects/deepgemm.cmake
new
+151
@@ -0,0 +1,151 @@
1
+include(FetchContent)
2
+
3
+# If DEEPGEMM_SRC_DIR is set, DeepGEMM is built from that directory
4
+# instead of downloading.
5
+# It can be set as an environment variable or passed as a cmake argument.
6
+# The environment variable takes precedence.
7
+if (DEFINED ENV{DEEPGEMM_SRC_DIR})
8
+ set(DEEPGEMM_SRC_DIR $ENV{DEEPGEMM_SRC_DIR})
9
+endif()
10
+
11
+if(DEEPGEMM_SRC_DIR)
12
+ FetchContent_Declare(
13
+ deepgemm
14
+ SOURCE_DIR ${DEEPGEMM_SRC_DIR}
15
+ CONFIGURE_COMMAND ""
16
+ BUILD_COMMAND ""
17
+ )
18
+else()
19
+ # This ref should be kept in sync with tools/install_deepgemm.sh
20
+ FetchContent_Declare(
21
+ deepgemm
22
+ GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
23
+ GIT_TAG 477618cd51baffca09c4b0b87e97c03fe827ef03
24
+ GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
25
+ GIT_PROGRESS TRUE
26
+ CONFIGURE_COMMAND ""
27
+ BUILD_COMMAND ""
28
+ )
29
+endif()
30
+
31
+# Use FetchContent_Populate (not MakeAvailable) to avoid processing
32
+# DeepGEMM's own CMakeLists.txt which has incompatible find_package calls.
33
+FetchContent_GetProperties(deepgemm)
34
+if(NOT deepgemm_POPULATED)
35
+ FetchContent_Populate(deepgemm)
36
+endif()
37
+message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}")
38
+
39
+# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100
40
+set(DEEPGEMM_SUPPORT_ARCHS)
41
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3)
42
+ list(APPEND DEEPGEMM_SUPPORT_ARCHS "9.0a")
43
+endif()
44
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
45
+ list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
46
+elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
47
+ list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
48
+endif()
49
+
50
+cuda_archs_loose_intersection(DEEPGEMM_ARCHS
51
+ "${DEEPGEMM_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
52
+
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>:-std=c++17>
80
+ $<$<COMPILE_LANGUAGE:CXX>:-O3>
81
+ $<$<COMPILE_LANGUAGE:CXX>:-Wno-psabi>
82
+ $<$<COMPILE_LANGUAGE:CXX>:-Wno-deprecated-declarations>)
83
+
84
+ # torch_python is required because DeepGEMM uses pybind11 type casters
85
+ # for at::Tensor (via PYBIND11_MODULE), unlike vLLM's own extensions which
86
+ # use torch::Library custom ops.
87
+ find_library(TORCH_PYTHON_LIBRARY torch_python
88
+ PATHS "${TORCH_INSTALL_PREFIX}/lib"
89
+ REQUIRED)
90
+
91
+ target_link_libraries(_deep_gemm_C PRIVATE
92
+ torch ${TORCH_LIBRARIES} "${TORCH_PYTHON_LIBRARY}"
93
+ CUDA::cudart CUDA::nvrtc)
94
+
95
+ # Install the shared library into the vendored package directory
96
+ install(TARGETS _deep_gemm_C
97
+ LIBRARY DESTINATION vllm/third_party/deep_gemm
98
+ COMPONENT _deep_gemm_C)
99
+
100
+ #
101
+ # Vendor DeepGEMM Python package files
102
+ #
103
+ install(FILES
104
+ "${deepgemm_SOURCE_DIR}/deep_gemm/__init__.py"
105
+ DESTINATION vllm/third_party/deep_gemm
106
+ COMPONENT _deep_gemm_C)
107
+
108
+ install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/utils/"
109
+ DESTINATION vllm/third_party/deep_gemm/utils
110
+ COMPONENT _deep_gemm_C
111
+ FILES_MATCHING PATTERN "*.py")
112
+
113
+ install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/testing/"
114
+ DESTINATION vllm/third_party/deep_gemm/testing
115
+ COMPONENT _deep_gemm_C
116
+ FILES_MATCHING PATTERN "*.py")
117
+
118
+ install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/legacy/"
119
+ DESTINATION vllm/third_party/deep_gemm/legacy
120
+ COMPONENT _deep_gemm_C
121
+ FILES_MATCHING PATTERN "*.py")
122
+
123
+ # Generate envs.py (normally generated by DeepGEMM's setup.py build step)
124
+ file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py"
125
+ "# Pre-installed environment variables\npersistent_envs = dict()\n")
126
+ install(FILES "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py"
127
+ DESTINATION vllm/third_party/deep_gemm
128
+ RENAME envs.py
129
+ COMPONENT _deep_gemm_C)
130
+
131
+ #
132
+ # Install include files needed for JIT compilation at runtime.
133
+ # The JIT compiler finds these relative to the package directory.
134
+ #
135
+
136
+ # DeepGEMM's own CUDA headers
137
+ install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/include/"
138
+ DESTINATION vllm/third_party/deep_gemm/include
139
+ COMPONENT _deep_gemm_C)
140
+
141
+ # CUTLASS and CuTe headers (vendored for JIT, separate from vLLM's CUTLASS)
142
+ install(DIRECTORY "${deepgemm_SOURCE_DIR}/third-party/cutlass/include/"
143
+ DESTINATION vllm/third_party/deep_gemm/include
144
+ COMPONENT _deep_gemm_C)
145
+
146
+else()
147
+ message(STATUS "DeepGEMM will not compile: "
148
+ "unsupported CUDA architecture ${CUDA_ARCHS}")
149
+ # Create empty target so setup.py doesn't fail on unsupported systems
150
+ add_custom_target(_deep_gemm_C)
151
+endif()
docker/Dockerfile
+1
-26
@@ -315,7 +315,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
315
#################### CSRC BUILD IMAGE ####################
316
317
#################### EXTENSIONS BUILD IMAGE ####################
318
-# Build DeepGEMM, DeepEP - runs in PARALLEL with csrc-build
318
+# Build DeepEP - runs in PARALLEL with csrc-build
319
# This stage is independent and doesn't affect csrc cache
320
FROM base AS extensions-build
321
ARG CUDA_VERSION
@@ -327,21 +327,6 @@ ENV UV_LINK_MODE=copy
327
328
WORKDIR /workspace
329
330
-# Build DeepGEMM wheel
331
-# Default moved here from tools/install_deepgemm.sh for centralized version management
332
-ARG DEEPGEMM_GIT_REF=477618cd51baffca09c4b0b87e97c03fe827ef03
333
-COPY tools/install_deepgemm.sh /tmp/install_deepgemm.sh
334
-RUN --mount=type=cache,target=/root/.cache/uv \
335
- mkdir -p /tmp/deepgemm/dist && \
336
- VLLM_DOCKER_BUILD_CONTEXT=1 TORCH_CUDA_ARCH_LIST="9.0a 10.0a" /tmp/install_deepgemm.sh \
337
- --cuda-version "${CUDA_VERSION}" \
338
- ${DEEPGEMM_GIT_REF:+--ref "$DEEPGEMM_GIT_REF"} \
339
- --wheel-dir /tmp/deepgemm/dist || \
340
- echo "DeepGEMM build skipped (CUDA version requirement not met)"
341
-
342
-# Ensure the wheel dir exists so COPY won't fail when DeepGEMM is skipped
343
-RUN mkdir -p /tmp/deepgemm/dist && touch /tmp/deepgemm/dist/.deepgemm_skipped
344
-
330
# Build DeepEP wheels
331
COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
332
# Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management
@@ -426,7 +411,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
411
python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
412
413
# Copy extension wheels from extensions-build stage for later use
429
-COPY --from=extensions-build /tmp/deepgemm/dist /tmp/deepgemm/dist
414
COPY --from=extensions-build /tmp/ep_kernels_workspace/dist /tmp/ep_kernels_workspace/dist
415
416
# Check the size of the wheel if RUN_WHEEL_CHECK is true
@@ -693,15 +677,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
677
. /etc/environment && \
678
uv pip list
679
696
-# Install deepgemm wheel that has been built in the `build` stage
697
-RUN --mount=type=cache,target=/root/.cache/uv \
698
- --mount=type=bind,from=build,source=/tmp/deepgemm/dist,target=/tmp/deepgemm/dist,ro \
699
- sh -c 'if ls /tmp/deepgemm/dist/*.whl >/dev/null 2>&1; then \
700
- uv pip install --system /tmp/deepgemm/dist/*.whl; \
701
- else \
702
- echo "No DeepGEMM wheels to install; skipping."; \
703
- fi'
704
-
680
# Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH
681
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
682
docker/versions.json
-3
@@ -52,9 +52,6 @@
52
"vllm_target_device": {
53
"default": "cuda"
54
},
55
- "DEEPGEMM_GIT_REF": {
56
- "default": "477618cd51baffca09c4b0b87e97c03fe827ef03"
57
- },
55
"DEEPEP_COMMIT_HASH": {
56
"default": "73b6ea4"
57
},
docs/assets/contributing/dockerfile-stages-dependency.png
Binary files a/docs/assets/contributing/dockerfile-stages-dependency.png and b/docs/assets/contributing/dockerfile-stages-dependency.png differ
setup.py
+29
@@ -379,6 +379,20 @@ class cmake_build_ext(build_ext):
379
dirs_exist_ok=True,
380
)
381
382
+ if _is_cuda():
383
+ # copy vendored deep_gemm package from build_lib to source tree
384
+ # for editable installs
385
+ deep_gemm_build = os.path.join(
386
+ self.build_lib, "vllm", "third_party", "deep_gemm"
387
+ )
388
+ if os.path.exists(deep_gemm_build):
389
+ print(f"Copying {deep_gemm_build} to vllm/third_party/deep_gemm")
390
+ shutil.copytree(
391
+ deep_gemm_build,
392
+ "vllm/third_party/deep_gemm",
393
+ dirs_exist_ok=True,
394
+ )
395
+
396
397
class precompiled_build_ext(build_ext):
398
"""Disables extension building when using precompiled binaries."""
@@ -685,6 +699,8 @@ class precompiled_wheel_utils:
699
flashmla_regex = re.compile(
700
r"vllm/third_party/flashmla/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
701
)
702
+ # DeepGEMM: extract all files (.py, .so, .cuh, .h, .hpp, etc.)
703
+ deep_gemm_regex = re.compile(r"vllm/third_party/deep_gemm/.*")
704
file_members = list(
705
filter(lambda x: x.filename in files_to_copy, wheel.filelist)
706
)
@@ -699,6 +715,9 @@ class precompiled_wheel_utils:
715
file_members += list(
716
filter(lambda x: flashmla_regex.match(x.filename), wheel.filelist)
717
)
718
+ file_members += list(
719
+ filter(lambda x: deep_gemm_regex.match(x.filename), wheel.filelist)
720
+ )
721
722
for file in file_members:
723
print(f"[extract] {file.filename}")
@@ -987,6 +1006,12 @@ if _is_cuda():
1006
ext_modules.append(
1007
CMakeExtension(name="vllm._flashmla_extension_C", optional=True)
1008
)
1009
+ if envs.VLLM_USE_PRECOMPILED or (
1010
+ CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3")
1011
+ ):
1012
+ # DeepGEMM requires CUDA 12.3+ (SM90/SM100)
1013
+ # Optional since it won't build on unsupported architectures
1014
+ ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True))
1015
1016
if _is_cpu():
1017
import platform
@@ -1014,6 +1039,10 @@ package_data = {
1039
"entrypoints/serve/instrumentator/static/*.js",
1040
"entrypoints/serve/instrumentator/static/*.css",
1041
"distributed/kv_transfer/kv_connector/v1/hf3fs/utils/*.cpp",
1042
+ # DeepGEMM JIT include headers (vendored via cmake)
1043
+ "third_party/deep_gemm/include/**/*.cuh",
1044
+ "third_party/deep_gemm/include/**/*.h",
1045
+ "third_party/deep_gemm/include/**/*.hpp",
1046
]
1047
}
1048
tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
+5
-3
@@ -14,7 +14,11 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
14
get_fp8_min_max,
15
)
16
from vllm.platforms import current_platform
17
-from vllm.utils.deep_gemm import DeepGemmQuantScaleFMT, has_deep_gemm
17
+from vllm.utils.deep_gemm import (
18
+ DeepGemmQuantScaleFMT,
19
+ has_deep_gemm,
20
+ transform_sf_into_required_layout,
21
+)
22
from vllm.utils.math_utils import cdiv, round_up
23
from vllm.utils.torch_utils import set_random_seed
24
@@ -256,8 +260,6 @@ def test_silu_mul_fp8_quant_deep_gemm(E: int, T: int, H: int, fp8_type: torch.dt
260
and current_platform.has_device_capability(100)
261
and scale_fmt == DeepGemmQuantScaleFMT.UE8M0
262
):
259
- from deep_gemm import transform_sf_into_required_layout
260
-
263
_q, _s = ref_with_scale_fmt(
264
E,
265
T,
tools/install_deepgemm.sh
+1
@@ -5,6 +5,7 @@
5
set -e
6
7
# Default values
8
+# Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake
9
DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
10
DEEPGEMM_GIT_REF="477618cd51baffca09c4b0b87e97c03fe827ef03"
11
WHEEL_DIR=""
vllm/utils/deep_gemm.py
+51
-3
@@ -138,6 +138,41 @@ _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None
138
_transform_sf_into_required_layout_impl: Callable[..., Any] | None = None
139
140
141
+def _import_deep_gemm():
142
+ """Import the deep_gemm module.
143
+
144
+ Prefers an externally installed ``deep_gemm`` package (so users can
145
+ pin a specific version), then falls back to the vendored copy bundled
146
+ in the vLLM wheel.
147
+
148
+ Returns ``None`` when neither source is usable.
149
+ """
150
+ # 1. Try the external (pip-installed) package first.
151
+ try:
152
+ module = importlib.import_module("deep_gemm")
153
+ logger.debug_once("Imported deep_gemm module from site-packages")
154
+ return module
155
+ except ImportError:
156
+ logger.debug_once(
157
+ "deep_gemm not found in site-packages, "
158
+ "trying vendored vllm.third_party.deep_gemm"
159
+ )
160
+
161
+ # 2. Fall back to the vendored copy bundled in the vLLM wheel.
162
+ try:
163
+ module = importlib.import_module("vllm.third_party.deep_gemm")
164
+ logger.debug_once("Imported deep_gemm module from vllm.third_party.deep_gemm")
165
+ return module
166
+ except ImportError:
167
+ logger.debug_once("Vendored deep_gemm not found either")
168
+ except Exception as e:
169
+ # The vendored module may raise RuntimeError during _C.init()
170
+ # if JIT include files are missing (e.g. incomplete wheel).
171
+ logger.warning_once("Failed to import vendored deep_gemm: %s", e)
172
+
173
+ return None
174
+
175
+
176
def _lazy_init() -> None:
177
"""Import deep_gemm and resolve symbols on first use."""
178
global _fp8_gemm_nt_impl, _grouped_impl, _grouped_masked_impl
@@ -169,7 +204,9 @@ def _lazy_init() -> None:
204
envs.VLLM_CACHE_ROOT, "deep_gemm"
205
)
206
172
- _dg = importlib.import_module("deep_gemm")
207
+ _dg = _import_deep_gemm()
208
+ if _dg is None:
209
+ return
210
211
_fp8_gemm_nt_impl = getattr(_dg, "fp8_gemm_nt", None)
212
_grouped_impl = getattr(_dg, "m_grouped_fp8_gemm_nt_contiguous", None)
@@ -193,8 +230,18 @@ def _lazy_init() -> None:
230
231
def get_num_sms() -> int:
232
_lazy_init()
196
- _dg = importlib.import_module("deep_gemm")
197
- return int(_dg.get_num_sms())
233
+ dg = _import_deep_gemm()
234
+ if dg is None:
235
+ raise RuntimeError("DeepGEMM is not available")
236
+ return int(dg.get_num_sms())
237
+
238
+
239
+def set_num_sms(num_sms: int) -> None:
240
+ _lazy_init()
241
+ dg = _import_deep_gemm()
242
+ if dg is None:
243
+ raise RuntimeError("DeepGEMM is not available")
244
+ dg.set_num_sms(num_sms)
245
246
247
@functools.cache
@@ -446,6 +493,7 @@ __all__ = [
493
"is_deep_gemm_e8m0_used",
494
"is_deep_gemm_supported",
495
"get_num_sms",
496
+ "set_num_sms",
497
"should_use_deepgemm_for_fp8_linear",
498
"get_col_major_tma_aligned_tensor",
499
"get_mk_alignment_for_contiguous_layout",
vllm/utils/import_utils.py
+7
-2
@@ -408,8 +408,13 @@ def has_deep_ep() -> bool:
408
409
410
def has_deep_gemm() -> bool:
411
- """Whether the optional `deep_gemm` package is available."""
412
- return _has_module("deep_gemm")
411
+ """Whether the optional `deep_gemm` package is available.
412
+
413
+ Prefers an externally installed ``deep_gemm`` package (so users can
414
+ override with a newer version), then falls back to the vendored copy
415
+ bundled in the vLLM wheel.
416
+ """
417
+ return _has_module("deep_gemm") or _has_module("vllm.third_party.deep_gemm")
418
419
420
def has_nixl_ep() -> bool:
vllm/v1/worker/gpu_ubatch_wrapper.py
+2
-3
@@ -23,6 +23,7 @@ from vllm.logger import init_logger
23
from vllm.model_executor.offloader.base import get_offloader
24
from vllm.platforms import current_platform
25
from vllm.sequence import IntermediateTensors
26
+from vllm.utils.deep_gemm import set_num_sms as deep_gemm_set_num_sms
27
from vllm.utils.import_utils import has_deep_gemm
28
from vllm.utils.platform_utils import num_compute_units
29
from vllm.v1.worker.ubatching import UBatchContext, make_ubatch_contexts
@@ -158,9 +159,7 @@ class UBatchWrapper:
159
# TODO(lucas): support other kernels besides DeepGEMM
160
set_compute_sms = lambda sms: None
161
if has_deep_gemm() and comm_sms > 0:
161
- import deep_gemm as dg
162
-
163
- set_compute_sms = lambda sms: dg.set_num_sms(sms)
162
+ set_compute_sms = lambda sms: deep_gemm_set_num_sms(sms)
163
164
return SMControlContextManager(
165
comm_sms=comm_sms,