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 14, 22:35 UTC
Add DeepGEMM to Dockerfile in vllm-base image (#21533)
Signed-off-by: Matthew Bonanni <mbonanni001@gmail.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Matthew Bonanni committed
Jul 31, 2025 at 21:01 UTC
e360316ab9902ecfc564710ae4b1539db867efd9
4 files changed
+46
-7
docker/Dockerfile
+28
-2
@@ -1,4 +1,3 @@
1
-
1
# The vLLM Dockerfile is used to construct vLLM image that can be directly used
2
# to run the OpenAI compatible server.
3
@@ -16,6 +15,7 @@ ARG PYTHON_VERSION=3.12
15
# Example:
16
# docker build --build-arg BUILD_BASE_IMAGE=registry.acme.org/mirror/nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04
17
ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04
18
+# TODO: Restore to base image after FlashInfer AOT wheel fixed
19
ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
20
21
# By parameterizing the Deadsnakes repository URL, we allow third-party to use
@@ -289,7 +289,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
289
290
#################### vLLM installation IMAGE ####################
291
# image with vLLM installed
292
-# TODO: Restore to base image after FlashInfer AOT wheel fixed
292
FROM ${FINAL_BASE_IMAGE} AS vllm-base
293
ARG CUDA_VERSION
294
ARG PYTHON_VERSION
@@ -435,6 +434,33 @@ RUN --mount=type=cache,target=/root/.cache/uv \
434
uv pip install --system -r requirements/build.txt \
435
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
436
437
+# Install DeepGEMM from source
438
+ARG DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
439
+ARG DEEPGEMM_GIT_REF="187656694f7f69e3e7975617a68bc3387680a7e1"
440
+RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'
441
+ . /etc/environment
442
+ CUDA_MAJOR="${CUDA_VERSION%%.*}"
443
+ CUDA_MINOR="${CUDA_VERSION#${CUDA_MAJOR}.}"
444
+ CUDA_MINOR="${CUDA_MINOR%%.*}"
445
+ if [ "$CUDA_MAJOR" -ge 12 ] && [ "$CUDA_MINOR" -ge 8 ]; then
446
+ git clone --recursive --shallow-submodules \
447
+ ${DEEPGEMM_GIT_REPO} deepgemm
448
+ echo "🏗️ Building DeepGEMM"
449
+ pushd deepgemm
450
+ git checkout ${DEEPGEMM_GIT_REF}
451
+ # Build DeepGEMM
452
+ # (Based on https://github.com/deepseek-ai/DeepGEMM/blob/main/install.sh)
453
+ rm -rf build dist
454
+ rm -rf *.egg-info
455
+ python3 setup.py bdist_wheel
456
+ uv pip install --system dist/*.whl
457
+ popd
458
+ rm -rf deepgemm
459
+ else
460
+ echo "Skipping DeepGEMM installation (requires CUDA 12.8+ but got ${CUDA_VERSION})"
461
+ fi
462
+BASH
463
+
464
#################### vLLM installation IMAGE ####################
465
466
#################### TEST IMAGE ####################
tests/kernels/moe/test_deepep_deepgemm_moe.py
+3
-2
@@ -20,7 +20,8 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
20
FusedMoEModularKernel)
21
from vllm.platforms import current_platform
22
from vllm.utils import has_deep_ep, has_deep_gemm
23
-from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used
23
+from vllm.utils.deep_gemm import (is_blackwell_deep_gemm_used,
24
+ is_deep_gemm_supported)
25
26
from .parallel_utils import ProcessGroupInfo, parallel_launch
27
from .utils import make_test_weights
@@ -46,7 +47,7 @@ requires_deep_ep = pytest.mark.skipif(
47
)
48
49
requires_deep_gemm = pytest.mark.skipif(
49
- not has_deep_gemm(),
50
+ not is_deep_gemm_supported(),
51
reason="Requires deep_gemm kernels",
52
)
53
tests/kernels/moe/test_deepgemm.py
+3
-3
@@ -15,13 +15,13 @@ import torch
15
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
16
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
17
per_token_group_quant_fp8)
18
-from vllm.utils import has_deep_gemm
19
-from vllm.utils.deep_gemm import calc_diff, per_block_cast_to_fp8
18
+from vllm.utils.deep_gemm import (calc_diff, is_deep_gemm_supported,
19
+ per_block_cast_to_fp8)
20
21
BLOCK_SIZE = [128, 128]
22
23
requires_deep_gemm = pytest.mark.skipif(
24
- not has_deep_gemm(),
24
+ not is_deep_gemm_supported(),
25
reason="Requires deep_gemm kernels",
26
)
27
vllm/utils/deep_gemm.py
+12
@@ -17,6 +17,17 @@ from vllm.platforms import current_platform
17
from vllm.utils import has_deep_gemm
18
19
20
+@functools.cache
21
+def is_deep_gemm_supported() -> bool:
22
+ """Return ``True`` if DeepGEMM is supported on the current platform.
23
+ Currently, only Hopper and Blackwell GPUs are supported.
24
+ """
25
+ supported_arch = current_platform.is_cuda() and (
26
+ current_platform.is_device_capability(90)
27
+ or current_platform.is_device_capability(100))
28
+ return has_deep_gemm() and supported_arch
29
+
30
+
31
@functools.cache
32
def is_blackwell_deep_gemm_used() -> bool:
33
"""Return ``True`` if vLLM is configured to use DeepGEMM on a
@@ -142,4 +153,5 @@ __all__ = [
153
"fp8_m_grouped_gemm_nt_masked",
154
"per_block_cast_to_fp8",
155
"is_blackwell_deep_gemm_used",
156
+ "is_deep_gemm_supported",
157
]