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 12, 02:15 UTC
[docker] Build CUDA kernels in separate Docker stage for faster rebuilds (#29452)
Signed-off-by: Amr Mahdi <amrmahdi@meta.com>
Amr Mahdi committed
Dec 3, 2025 at 03:41 UTC
f5d3d93c40417c296c20dc301100e55708a17f3f
4 files changed
+74
-11
docker/Dockerfile
+58
-8
@@ -150,8 +150,8 @@ ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
150
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
151
#################### BASE BUILD IMAGE ####################
152
153
-#################### WHEEL BUILD IMAGE ####################
154
-FROM base AS build
153
+#################### CSRC BUILD IMAGE ####################
154
+FROM base AS csrc-build
155
ARG TARGETPLATFORM
156
157
ARG PIP_INDEX_URL UV_INDEX_URL
@@ -172,10 +172,13 @@ RUN --mount=type=cache,target=/root/.cache/uv \
172
uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \
173
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
174
175
-COPY . .
176
-ARG GIT_REPO_CHECK=0
177
-RUN --mount=type=bind,source=.git,target=.git \
178
- if [ "$GIT_REPO_CHECK" != "0" ]; then bash tools/check_repo.sh ; fi
175
+WORKDIR /workspace
176
+
177
+COPY pyproject.toml setup.py CMakeLists.txt ./
178
+COPY cmake cmake/
179
+COPY csrc csrc/
180
+COPY vllm/envs.py vllm/envs.py
181
+COPY vllm/__init__.py vllm/__init__.py
182
183
# max jobs used by Ninja to build extensions
184
ARG max_jobs=2
@@ -195,9 +198,11 @@ ARG SCCACHE_S3_NO_CREDENTIALS=0
198
ARG VLLM_USE_PRECOMPILED=""
199
ARG VLLM_MAIN_CUDA_VERSION=""
200
201
+# Use dummy version for csrc-build wheel (only .so files are extracted, version doesn't matter)
202
+ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+csrc.build"
203
+
204
# if USE_SCCACHE is set, use sccache to speed up compilation
205
RUN --mount=type=cache,target=/root/.cache/uv \
200
- --mount=type=bind,source=.git,target=.git \
206
if [ "$USE_SCCACHE" = "1" ]; then \
207
echo "Installing sccache..." \
208
&& curl -L -o sccache.tar.gz ${SCCACHE_DOWNLOAD_URL} \
@@ -223,7 +228,6 @@ ENV VLLM_TARGET_DEVICE=${vllm_target_device}
228
ENV CCACHE_DIR=/root/.cache/ccache
229
RUN --mount=type=cache,target=/root/.cache/ccache \
230
--mount=type=cache,target=/root/.cache/uv \
226
- --mount=type=bind,source=.git,target=.git \
231
if [ "$USE_SCCACHE" != "1" ]; then \
232
# Clean any existing CMake artifacts
233
rm -rf .deps && \
@@ -232,6 +236,52 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
236
export VLLM_DOCKER_BUILD_CONTEXT=1 && \
237
python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
238
fi
239
+#################### CSRC BUILD IMAGE ####################
240
+
241
+#################### WHEEL BUILD IMAGE ####################
242
+FROM base AS build
243
+ARG TARGETPLATFORM
244
+
245
+ARG PIP_INDEX_URL UV_INDEX_URL
246
+ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
247
+ARG PYTORCH_CUDA_INDEX_BASE_URL
248
+
249
+# install build dependencies
250
+COPY requirements/build.txt requirements/build.txt
251
+
252
+# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
253
+# Reference: https://github.com/astral-sh/uv/pull/1694
254
+ENV UV_HTTP_TIMEOUT=500
255
+ENV UV_INDEX_STRATEGY="unsafe-best-match"
256
+# Use copy mode to avoid hardlink failures with Docker cache mounts
257
+ENV UV_LINK_MODE=copy
258
+
259
+RUN --mount=type=cache,target=/root/.cache/uv \
260
+ uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \
261
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
262
+
263
+WORKDIR /workspace
264
+
265
+COPY --from=csrc-build /workspace/dist /precompiled-wheels
266
+
267
+COPY . .
268
+
269
+ARG GIT_REPO_CHECK=0
270
+RUN --mount=type=bind,source=.git,target=.git \
271
+ if [ "$GIT_REPO_CHECK" != "0" ]; then bash tools/check_repo.sh ; fi
272
+
273
+ARG vllm_target_device="cuda"
274
+ENV VLLM_TARGET_DEVICE=${vllm_target_device}
275
+
276
+# Skip adding +precompiled suffix to version (preserves git-derived version)
277
+ENV VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX=1
278
+
279
+RUN --mount=type=cache,target=/root/.cache/uv \
280
+ --mount=type=bind,source=.git,target=.git \
281
+ if [ "${vllm_target_device}" = "cuda" ]; then \
282
+ export VLLM_PRECOMPILED_WHEEL_LOCATION=$(ls /precompiled-wheels/*.whl); \
283
+ fi && \
284
+ python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
285
286
# Install DeepGEMM from source
287
ARG DEEPGEMM_GIT_REF
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
+11
-3
@@ -461,14 +461,22 @@ class precompiled_wheel_utils:
461
"vllm/cumem_allocator.abi3.so",
462
]
463
464
- compiled_regex = re.compile(
464
+ flash_attn_regex = re.compile(
465
r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
466
)
467
+ triton_kernels_regex = re.compile(
468
+ r"vllm/third_party/triton_kernels/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
469
+ )
470
file_members = list(
471
filter(lambda x: x.filename in files_to_copy, wheel.filelist)
472
)
473
file_members += list(
471
- filter(lambda x: compiled_regex.match(x.filename), wheel.filelist)
474
+ filter(lambda x: flash_attn_regex.match(x.filename), wheel.filelist)
475
+ )
476
+ file_members += list(
477
+ filter(
478
+ lambda x: triton_kernels_regex.match(x.filename), wheel.filelist
479
+ )
480
)
481
482
for file in file_members:
@@ -648,7 +656,7 @@ def get_vllm_version() -> str:
656
if envs.VLLM_TARGET_DEVICE == "empty":
657
version += f"{sep}empty"
658
elif _is_cuda():
651
- if envs.VLLM_USE_PRECOMPILED:
659
+ if envs.VLLM_USE_PRECOMPILED and not envs.VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX:
660
version += f"{sep}precompiled"
661
else:
662
cuda_version = str(get_nvcc_cuda_version())
vllm/envs.py
+5
@@ -78,6 +78,7 @@ if TYPE_CHECKING:
78
MAX_JOBS: str | None = None
79
NVCC_THREADS: str | None = None
80
VLLM_USE_PRECOMPILED: bool = False
81
+ VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX: bool = False
82
VLLM_DOCKER_BUILD_CONTEXT: bool = False
83
VLLM_TEST_USE_PRECOMPILED_NIGHTLY_WHEEL: bool = False
84
VLLM_KEEP_ALIVE_ON_ENGINE_DEATH: bool = False
@@ -462,6 +463,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
463
.lower()
464
in ("1", "true")
465
or bool(os.environ.get("VLLM_PRECOMPILED_WHEEL_LOCATION")),
466
+ # If set, skip adding +precompiled suffix to version string
467
+ "VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX": lambda: bool(
468
+ int(os.environ.get("VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX", "0"))
469
+ ),
470
# Used to mark that setup.py is running in a Docker build context,
471
# in order to force the use of precompiled binaries.
472
"VLLM_DOCKER_BUILD_CONTEXT": lambda: os.environ.get("VLLM_DOCKER_BUILD_CONTEXT", "")