@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3593c964de

[ROCm] Add TheRock preview docker updates, Keep Python 3.12 and Ubuntu 22.04 (#49925)

Signed-off-by: Randall Smith <Randall.Smith@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

rasmith committed Aug 31, 2026 at 11:09 UTC 3593c964de198646a2a284c24fc53cccd21afd7e
7 files changed +1384 -7
docker/Dockerfile.rock new
+992
@@ -0,0 +1,992 @@
1 +# default base image
2 +ARG REMOTE_VLLM="0"
3 +ARG VLLM_REPO="https://github.com/vllm-project/vllm.git"
4 +ARG VLLM_BRANCH="main"
5 +ARG COMMON_WORKDIR=/app
6 +ARG BASE_IMAGE=rocm/vllm-dev:base
7 +ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base
8 +ARG ROCM_TRITON_KERNELS_COMMIT=0f380657dbf3ee86eb57558ff71df24f03b5d4e7
9 +ARG INSTALL_LMCACHE=false
10 +# NIC backend for MoRI RDMA support.
11 +# By default (all), drivers and userspace libraries for all supported NIC types
12 +# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime.
13 +# To install drivers for a single NIC type only, set NIC_BACKEND explicitly:
14 +# --build-arg NIC_BACKEND=ainic # AMD AINIC (Pensando) only
15 +# --build-arg NIC_BACKEND=bnxt # Broadcom Thor-2 only
16 +# --build-arg NIC_BACKEND=none # Install nothing.
17 +ARG NIC_BACKEND=all
18 +# AMD AINIC apt repo settings
19 +# Users can specify a custom version compatible with their host drivers.
20 +# The default version has been tested with ioinic-dkms=25.11.1.001
21 +ARG AINIC_VERSION=1.117.3-hydra
22 +ARG UBUNTU_CODENAME=jammy
23 +
24 +# Sccache configuration. Release builds use this today; CI can opt in when a
25 +# shared S3-compatible cache backend is available.
26 +ARG USE_SCCACHE
27 +ARG SCCACHE_DOWNLOAD_URL
28 +ARG SCCACHE_ENDPOINT
29 +ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
30 +ARG SCCACHE_REGION_NAME=us-west-2
31 +ARG SCCACHE_S3_NO_CREDENTIALS=0
32 +
33 +FROM ${BASE_IMAGE} AS base
34 +
35 +ARG ARG_PYTORCH_ROCM_ARCH
36 +ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
37 +
38 +# Install build dependencies and utilities
39 +RUN apt-get update -q -y && apt-get install -q -y \
40 + sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \
41 + apt-transport-https ca-certificates wget curl \
42 + libnuma-dev ccache mold
43 +RUN --mount=type=cache,target=/root/.cache/pip \
44 + python3 -m pip install --upgrade pip
45 +# Note: mold is installed but not set as the system default linker because
46 +# some packages use JIT compilation at runtime with flags mold does not support.
47 +# Build stages opt in via LDFLAGS="-fuse-ld=mold".
48 +# Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base)
49 +ARG USE_SCCACHE
50 +RUN if [ "$USE_SCCACHE" != "1" ]; then \
51 + apt-get purge -y sccache || true; \
52 + python3 -m pip uninstall -y sccache || true; \
53 + rm -f "$(which sccache)" || true; \
54 + fi
55 +
56 +# Install UV — download first, then run, so a curl failure is not masked by the pipe
57 +RUN curl -LsSf --retry 3 --retry-delay 5 https://astral.sh/uv/install.sh -o /tmp/uv-install.sh \
58 + && env UV_INSTALL_DIR="/usr/local/bin" sh /tmp/uv-install.sh \
59 + && rm -f /tmp/uv-install.sh \
60 + && uv --version
61 +
62 +# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
63 +# Reference: https://github.com/astral-sh/uv/pull/1694
64 +ENV UV_HTTP_TIMEOUT=500
65 +ENV UV_INDEX_STRATEGY="unsafe-best-match"
66 +# Use copy mode to avoid hardlink failures with Docker cache mounts
67 +ENV UV_LINK_MODE=copy
68 +# ccache directory - persisted across layer rebuilds via cache mounts.
69 +ENV CCACHE_DIR=/root/.cache/ccache
70 +ENV CCACHE_COMPILERCHECK=content
71 +# Empty by default so build steps fall back to $(nproc); CI can override.
72 +ARG max_jobs
73 +ENV MAX_JOBS=${max_jobs}
74 +
75 +# Install sccache if USE_SCCACHE is enabled (for release builds)
76 +ARG USE_SCCACHE
77 +ARG SCCACHE_DOWNLOAD_URL
78 +ARG SCCACHE_ENDPOINT
79 +ARG SCCACHE_BUCKET_NAME
80 +ARG SCCACHE_REGION_NAME
81 +ARG SCCACHE_S3_NO_CREDENTIALS
82 +RUN if [ "$USE_SCCACHE" = "1" ]; then \
83 + if command -v sccache >/dev/null 2>&1; then \
84 + echo "sccache already installed, skipping installation"; \
85 + sccache --version; \
86 + else \
87 + echo "Installing sccache..." \
88 + && SCCACHE_ARCH="x86_64" \
89 + && SCCACHE_VERSION="v0.8.1" \
90 + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \
91 + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \
92 + && tar -xzf /tmp/sccache.tar.gz -C /tmp \
93 + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
94 + && chmod +x /usr/bin/sccache \
95 + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \
96 + && sccache --version; \
97 + fi; \
98 + fi
99 +
100 +# Set sccache environment variables only when USE_SCCACHE=1
101 +# This prevents S3 config from leaking into images when sccache is not used
102 +ARG USE_SCCACHE
103 +ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
104 +ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
105 +ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
106 +ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
107 +
108 +ARG COMMON_WORKDIR
109 +WORKDIR ${COMMON_WORKDIR}
110 +
111 +
112 +# -----------------------
113 +# vLLM fetch stages
114 +FROM base AS fetch_vllm_0
115 +ONBUILD COPY ./ vllm/
116 +FROM base AS fetch_vllm_1
117 +ARG VLLM_REPO
118 +ARG VLLM_BRANCH
119 +ONBUILD RUN git clone ${VLLM_REPO} \
120 + && cd vllm \
121 + && git fetch -v --prune -- origin ${VLLM_BRANCH} \
122 + && git checkout FETCH_HEAD \
123 + && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \
124 + git remote add upstream "https://github.com/vllm-project/vllm.git" \
125 + && git fetch upstream ; fi
126 +FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm
127 +
128 +# Resolve the device-independent source version once for every downstream
129 +# artifact. setuptools-scm reads Git metadata for local/remote sources and the
130 +# archive metadata generated for canonical CI contexts.
131 +FROM fetch_vllm AS vllm-version
132 +ARG COMMON_WORKDIR
133 +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
134 + cd ${COMMON_WORKDIR}/vllm \
135 + && uv pip install --system -r requirements/build/rust.txt \
136 + && python3 -c \
137 + "from setuptools_scm import get_version; print(get_version())" \
138 + > /vllm-version.txt
139 +
140 +# -----------------------
141 +# Rust toolchain inputs
142 +#
143 +# Keep these separate from Rust source inputs so Rust source-only changes do not
144 +# invalidate the long-lived ci_base image that carries the toolchain.
145 +FROM scratch AS rust_toolchain_input_0
146 +COPY rust-toolchain.toml /rust-toolchain-input/vllm/rust-toolchain.toml
147 +
148 +FROM base AS rust_toolchain_input_1
149 +ARG VLLM_REPO
150 +ARG VLLM_BRANCH
151 +RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-toolchain-input/vllm \
152 + && cd /rust-toolchain-input/vllm \
153 + && git fetch -v --prune -- origin ${VLLM_BRANCH} \
154 + && git sparse-checkout init --no-cone \
155 + && git sparse-checkout set \
156 + rust-toolchain.toml \
157 + && git checkout FETCH_HEAD
158 +
159 +FROM rust_toolchain_input_${REMOTE_VLLM} AS rust-toolchain-input
160 +
161 +# -----------------------
162 +# Rust build inputs
163 +#
164 +# REMOTE_VLLM=0: use only the local files that can affect Rust artifacts.
165 +# REMOTE_VLLM=1: fetch only those same files from the requested remote ref.
166 +FROM scratch AS rust_input_0
167 +COPY requirements/build/rust.txt /rust-input/vllm/requirements/build/rust.txt
168 +COPY rust/Cargo.lock rust/Cargo.toml /rust-input/vllm/rust/
169 +COPY rust/proto /rust-input/vllm/rust/proto
170 +COPY rust/src /rust-input/vllm/rust/src
171 +COPY rust-toolchain.toml /rust-input/vllm/rust-toolchain.toml
172 +COPY tools/build_rust.py /rust-input/vllm/tools/build_rust.py
173 +COPY build_rust.sh /rust-input/vllm/build_rust.sh
174 +
175 +FROM base AS rust_input_1
176 +ARG VLLM_REPO
177 +ARG VLLM_BRANCH
178 +RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-input/vllm \
179 + && cd /rust-input/vllm \
180 + && git fetch -v --prune -- origin ${VLLM_BRANCH} \
181 + && git sparse-checkout init --no-cone \
182 + && git sparse-checkout set \
183 + requirements/build/rust.txt \
184 + rust/Cargo.lock \
185 + rust/Cargo.toml \
186 + rust/proto/** \
187 + rust/src/** \
188 + rust-toolchain.toml \
189 + tools/build_rust.py \
190 + build_rust.sh \
191 + && git checkout FETCH_HEAD
192 +
193 +FROM rust_input_${REMOTE_VLLM} AS rust-input
194 +
195 +# -----------------------
196 +# Rust toolchain
197 +FROM base AS rust-toolchain
198 +
199 +ENV CARGO_HOME=/root/.cargo
200 +ENV RUSTUP_HOME=/root/.rustup
201 +ENV PATH=${CARGO_HOME}/bin:${PATH}
202 +
203 +RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
204 + ca-certificates curl \
205 + && rm -rf /var/lib/apt/lists/*
206 +
207 +COPY --from=rust-toolchain-input /rust-toolchain-input/vllm/rust-toolchain.toml /tmp/rust-toolchain.toml
208 +RUN TOOLCHAIN="$(grep '^channel' /tmp/rust-toolchain.toml | sed 's/.*= *"\(.*\)"/\1/')" \
209 + && if ! command -v rustup >/dev/null 2>&1; then \
210 + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \
211 + | sh -s -- -y --default-toolchain none; \
212 + fi \
213 + && rustup toolchain install "${TOOLCHAIN}" \
214 + && rustup default "${TOOLCHAIN}" \
215 + && rustc --version \
216 + && cargo --version
217 +
218 +# -----------------------
219 +# Rust build stage
220 +# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages
221 +# don't need the rust toolchain.
222 +FROM rust-toolchain AS rust-build
223 +ARG COMMON_WORKDIR
224 +ARG USE_SCCACHE
225 +
226 +# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit
227 +# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
228 +ENV CARGO_BUILD_JOBS=4
229 +ENV CARGO_NET_RETRY=10
230 +ENV RUSTUP_MAX_RETRIES=10
231 +
232 +WORKDIR ${COMMON_WORKDIR}/vllm
233 +
234 +COPY --from=rust-input /rust-input/vllm/requirements/build/rust.txt requirements/build/rust.txt
235 +
236 +# BuildKit can run this stage in parallel with ROCm native builds. Keep Rust on
237 +# a separate local sccache daemon while sharing the same remote cache backend.
238 +ENV SCCACHE_SERVER_PORT=4227
239 +
240 +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
241 + uv pip install --system -r requirements/build/rust.txt
242 +
243 +COPY --from=rust-input /rust-input/vllm/rust rust
244 +COPY --from=rust-input /rust-input/vllm/rust-toolchain.toml rust-toolchain.toml
245 +COPY --from=rust-input /rust-input/vllm/tools/build_rust.py tools/build_rust.py
246 +COPY --from=rust-input /rust-input/vllm/build_rust.sh build_rust.sh
247 +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt
248 +
249 +# Build the release binary. Cargo's registry/git caches can be written by
250 +# concurrent BuildKit jobs on shared workers, so lock those cache mounts while
251 +# keeping the cache benefit. Do not cache target/, because stale target metadata
252 +# can outlive source updates across BuildKit cache reuse.
253 +RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
254 + --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
255 + if [ "$USE_SCCACHE" = "1" ]; then \
256 + export RUSTC_WRAPPER=sccache \
257 + && sccache --show-stats; \
258 + fi \
259 + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \
260 + && bash build_rust.sh \
261 + && test -x vllm/vllm-rs \
262 + && if [ "$USE_SCCACHE" = "1" ]; then \
263 + sccache --show-stats; \
264 + fi
265 +
266 +# -----------------------
267 +# Shared vLLM build dependencies
268 +#
269 +# Install requirements once, before bringing source-specific inputs into either
270 +# consumer. Native compilation and wheel packaging then share this expensive
271 +# layer while retaining their independent source cache boundaries.
272 +FROM base AS build_vllm_dependencies
273 +ARG COMMON_WORKDIR
274 +ENV VLLM_TARGET_DEVICE=rocm
275 +WORKDIR ${COMMON_WORKDIR}/vllm
276 +
277 +COPY --from=fetch_vllm \
278 + ${COMMON_WORKDIR}/vllm/requirements/common.txt \
279 + ${COMMON_WORKDIR}/vllm/requirements/rocm.txt \
280 + requirements/
281 +
282 +# Use constraints file to protect triton, pybind11, and torch dependencies.
283 +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
284 + uv pip install --system -c /etc/rocm-constraints.txt -r requirements/rocm.txt
285 +
286 +# Fetch only the pinned ROCm Triton kernels package. Keeping this before the
287 +# vLLM source COPYs makes the download reusable when native sources change.
288 +FROM base AS rocm-triton-kernels
289 +ARG ROCM_TRITON_KERNELS_COMMIT
290 +ADD https://codeload.github.com/ROCm/triton/tar.gz/${ROCM_TRITON_KERNELS_COMMIT} \
291 + /tmp/rocm-triton.tar.gz
292 +RUN mkdir -p /opt/rocm-triton-kernels \
293 + && tar --extract --gzip --file /tmp/rocm-triton.tar.gz \
294 + --directory /opt/rocm-triton-kernels --strip-components=4 \
295 + --no-same-owner \
296 + "triton-${ROCM_TRITON_KERNELS_COMMIT}/python/triton_kernels/triton_kernels" \
297 + && test -f /opt/rocm-triton-kernels/__init__.py \
298 + && rm /tmp/rocm-triton.tar.gz
299 +
300 +# -----------------------
301 +# vLLM native build stages
302 +#
303 +# csrc-build intentionally copies only files that affect ROCm native extension
304 +# compilation. That keeps unrelated CI/test/docs edits from invalidating the
305 +# expensive HIP/C++ build layer.
306 +FROM build_vllm_dependencies AS csrc-build
307 +ARG ROCM_TRITON_KERNELS_COMMIT
308 +COPY --from=rocm-triton-kernels \
309 + /opt/rocm-triton-kernels /opt/rocm-triton-kernels
310 +ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels
311 +
312 +# pyproject.toml is bind-mounted so it is not persisted in the csrc-build image,
313 +# but it still participates in the BuildKit cache key for this RUN step.
314 +COPY setup.py CMakeLists.txt ./
315 +COPY tools/build_rust.py tools/build_rust.py
316 +COPY cmake cmake/
317 +COPY csrc csrc/
318 +COPY vllm/envs.py vllm/envs.py
319 +COPY vllm/__init__.py vllm/__init__.py
320 +
321 +ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build"
322 +
323 +RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \
324 + --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
325 + grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \
326 + cmake/external_projects/triton_kernels.cmake \
327 + && export CCACHE_BASEDIR="$PWD" \
328 + && echo "=== ccache stats before ROCm native build ===" \
329 + && (ccache --show-stats || true) \
330 + && (ccache --zero-stats || true) \
331 + && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \
332 + && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \
333 + && PYTHONDONTWRITEBYTECODE=1 LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \
334 + && test -d dist \
335 + && ls dist/*.whl >/dev/null \
336 + && echo "=== ccache stats after ROCm native build ===" \
337 + && (ccache --show-stats || true)
338 +
339 +# Build the full vLLM ROCm wheel by reusing the native extension wheel from
340 +# csrc-build. This stage still rebuilds for Python/package changes, but skips
341 +# the expensive dependency installation and HIP/C++ compilation when their
342 +# respective inputs are unchanged.
343 +FROM build_vllm_dependencies AS build_vllm
344 +COPY --from=fetch_vllm ${COMMON_WORKDIR}/vllm ${COMMON_WORKDIR}/vllm
345 +COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels
346 +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt
347 +
348 +# Drop the pre-built Rust artifacts into the source tree. setup.py detects
349 +# them and ships them as-is, skipping the local Rust build.
350 +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
351 +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
352 +
353 +RUN cd "${COMMON_WORKDIR}/vllm" \
354 + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \
355 + && export VLLM_USE_PRECOMPILED=1 \
356 + && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \
357 + && export VLLM_DOCKER_BUILD_CONTEXT=1 \
358 + && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \
359 + && PYTHONDONTWRITEBYTECODE=1 python3 setup.py bdist_wheel --dist-dir=dist \
360 + && test -d dist \
361 + && ls dist/*.whl >/dev/null
362 +FROM scratch AS export_vllm
363 +ARG COMMON_WORKDIR
364 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl /
365 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements
366 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
367 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests
368 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples
369 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/scripts/benchmark_helion_kernels.py /scripts/benchmark_helion_kernels.py
370 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
371 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile
372 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu
373 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
374 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base
375 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl
376 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl
377 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl
378 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
379 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
380 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust
381 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
382 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
383 +
384 +# NIXL/UCX build stages
385 +FROM base AS build_nixl
386 +ARG NIXL_BRANCH="231d56753047c989062a5cb2ac703a1ad761c7d2"
387 +ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
388 +ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
389 +ARG UCX_REPO="https://github.com/openucx/ucx.git"
390 +ENV UCX_HOME=/usr/local/ucx
391 +ENV NIXL_HOME=/usr/local/nixl
392 +
393 +# NIXL build system dependencies and RDMA support
394 +RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
395 + libgrpc-dev \
396 + libgrpc++-dev \
397 + libprotobuf-dev \
398 + protobuf-compiler-grpc \
399 + libcpprest-dev \
400 + libaio-dev \
401 + librdmacm1 \
402 + librdmacm-dev \
403 + libibverbs1 \
404 + libibverbs-dev \
405 + ibverbs-utils \
406 + rdmacm-utils \
407 + ibverbs-providers \
408 + && rm -rf /var/lib/apt/lists/*
409 +
410 +RUN --mount=type=cache,target=/root/.cache/uv \
411 + uv pip install --system meson meson-python pyyaml types-PyYAML \
412 + auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0"
413 +
414 +RUN --mount=type=cache,target=/root/.cache/ccache \
415 + cd /usr/local/src && \
416 + git clone ${UCX_REPO} && \
417 + cd ucx && \
418 + git checkout ${UCX_BRANCH} && \
419 + ./autogen.sh && \
420 + mkdir build && cd build && \
421 + CC="ccache gcc" CXX="ccache g++" \
422 + ../configure \
423 + --prefix=/usr/local/ucx \
424 + --enable-shared \
425 + --disable-static \
426 + --disable-doxygen-doc \
427 + --enable-optimizations \
428 + --enable-devel-headers \
429 + --with-rocm=${ROCM_PATH} \
430 + --with-verbs \
431 + --with-dm \
432 + --enable-mt && \
433 + make -j$(nproc) && \
434 + make install
435 +
436 +ENV PATH=/usr/local/ucx/bin:$PATH
437 +ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
438 +
439 +RUN --mount=type=cache,target=/root/.cache/ccache \
440 + git clone ${NIXL_REPO} /opt/nixl && \
441 + cd /opt/nixl && \
442 + git checkout ${NIXL_BRANCH} && \
443 + CC="ccache gcc" CXX="ccache g++" \
444 + meson setup build --prefix=${NIXL_HOME} \
445 + -Ducx_path=${UCX_HOME} \
446 + -Dwheel_variant=rocm \
447 + -Dbuild_tests=false \
448 + -Dbuild_examples=false && \
449 + cd build && \
450 + ninja -j$(nproc) && \
451 + ninja install && \
452 + echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu" \
453 + > /etc/ld.so.conf.d/nixl.conf && \
454 + echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins" \
455 + >> /etc/ld.so.conf.d/nixl.conf && \
456 + ldconfig
457 +
458 +# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA,
459 +# so configure the ROCm wheel variant directly through Meson.
460 +# Exclude libcore and libpull from auditwheel: transitive dependencies
461 +# that are not shipped in the wheel and vary across base images.
462 +RUN cd /opt/nixl && \
463 + ./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \
464 + CC="ccache gcc" CXX="ccache g++" \
465 + uv build --wheel --no-build-isolation --out-dir /tmp/nixl_wheels \
466 + --python ${PYTHON_VERSION} \
467 + -Csetup-args=-Ducx_path=${UCX_HOME} \
468 + -Csetup-args=-Dwheel_variant=rocm \
469 + -Csetup-args=-Dbuild_tests=false \
470 + -Csetup-args=-Dbuild_examples=false && \
471 + mkdir -p /tmp/nixl_wheels/repaired /app/install && \
472 + auditwheel repair \
473 + --exclude 'libamdhip64*' \
474 + --exclude 'libcore*' \
475 + --exclude 'libpull*' \
476 + /tmp/nixl_wheels/nixl_rocm*.whl \
477 + --plat manylinux_2_34_$(uname -m) \
478 + --wheel-dir /tmp/nixl_wheels/repaired && \
479 + ./contrib/wheel_add_ucx_plugins.py \
480 + --ucx-plugins-dir ${UCX_HOME}/lib/ucx \
481 + --nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \
482 + /tmp/nixl_wheels/repaired/*.whl && \
483 + cp /tmp/nixl_wheels/repaired/*.whl /app/install
484 +
485 +# Fetch and verify pinned LMCache source.
486 +FROM base AS lmcache_source
487 +ARG LMCACHE_REPO="https://github.com/LMCache/LMCache.git"
488 +ARG LMCACHE_REF="140819c9d57a975dbc5678a6459a218e544cb58b"
489 +ARG LMCACHE_VERSION="0.5.3"
490 +
491 +# Remove NVIDIA-only dependencies and accept the release image's NumPy version.
492 +RUN printf '%s\n' "${LMCACHE_REF}" | grep -Eq '^[0-9a-f]{40}$' \
493 + && git clone --depth=1 --branch "v${LMCACHE_VERSION}" \
494 + "${LMCACHE_REPO}" /app/lmcache \
495 + && test "$(git -C /app/lmcache rev-parse HEAD)" = "${LMCACHE_REF}" \
496 + && test "$(grep -Ec '^(cufile-python|nvtx)([<=>[:space:]]|$)' \
497 + /app/lmcache/requirements/common.txt)" -eq 2 \
498 + && test "$(grep -Fxc 'numpy<=2.2.6' \
499 + /app/lmcache/requirements/common.txt)" -eq 1 \
500 + && sed -i -E \
501 + -e '/^(cufile-python|nvtx)([<=>[:space:]]|$)/d' \
502 + -e 's/^numpy<=2\.2\.6$/numpy<=2.3.5/' \
503 + /app/lmcache/requirements/common.txt \
504 + && rm -rf /app/lmcache/.git
505 +
506 +FROM base AS build_lmcache
507 +COPY --from=lmcache_source \
508 + /app/lmcache/requirements/build.txt /tmp/lmcache-build-requirements.txt
509 +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
510 + uv pip install --system -r /tmp/lmcache-build-requirements.txt
511 +
512 +ARG LMCACHE_VERSION="0.5.3"
513 +ARG LMCACHE_ROCM_ARCH="gfx942;gfx950"
514 +ARG USE_SCCACHE
515 +
516 +COPY --from=lmcache_source /app/lmcache /app/lmcache
517 +WORKDIR /app/lmcache
518 +
519 +RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
520 + set -eu; \
521 + printf '%s\n' "${LMCACHE_ROCM_ARCH}" \
522 + | grep -Eq '^gfx[0-9a-z]+(;gfx[0-9a-z]+)*$'; \
523 + offload_flags="$(printf '%s' "${LMCACHE_ROCM_ARCH}" \
524 + | sed 's/;/ --offload-arch=/g; s/^/--offload-arch=/')"; \
525 + if [ "${USE_SCCACHE:-0}" = "1" ]; then \
526 + export HIP_CLANG_PATH=/opt/sccache-wrappers; \
527 + fi; \
528 + mkdir -p /app/install; \
529 + SETUPTOOLS_SCM_PRETEND_VERSION="${LMCACHE_VERSION}" \
530 + PYTORCH_ROCM_ARCH="${LMCACHE_ROCM_ARCH}" \
531 + HIPCC_COMPILE_FLAGS_APPEND="${offload_flags}" \
532 + BUILD_WITH_HIP=1 CXX=hipcc \
533 + python3 setup.py bdist_wheel --dist-dir=/app/install; \
534 + test "$(find /app/install -maxdepth 1 -name '*.whl' | wc -l)" -eq 1
535 +
536 +RUN --mount=type=tmpfs,target=/tmp/lmcache-wheel \
537 + set -eu; \
538 + python3 -m zipfile -e /app/install/*.whl /tmp/lmcache-wheel; \
539 + c_ops="$(find /tmp/lmcache-wheel -type f -name 'c_ops*.so')"; \
540 + test "$(find /tmp/lmcache-wheel -type f -name 'c_ops*.so' | wc -l)" -eq 1; \
541 + expected="$(printf '%s' "${LMCACHE_ROCM_ARCH}" | tr ';' '\n' | sort -u)"; \
542 + actual="$(llvm-objdump --offloading "${c_ops}" \
543 + | grep -oE 'gfx[0-9a-z]+' | sort -u)"; \
544 + if [ "${actual}" != "${expected}" ]; then \
545 + printf 'LMCache code objects mismatch\nexpected:\n%s\nactual:\n%s\n' \
546 + "${expected}" "${actual}" >&2; \
547 + exit 1; \
548 + fi
549 +
550 +# DeepEP build stage. ROCShmem is provided by TheRock SDK
551 +# ROCSHMEM_DIR set to the SDK install path so DeepEP finds the headers, etc.
552 +FROM base AS build_deepep
553 +ENV ROCSHMEM_DIR=${ROCM_PATH}
554 +ARG DEEPEP_BRANCH="0f63d3e9"
555 +ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
556 +ARG DEEPEP_NIC="cx7"
557 +
558 +# Build DeepEP wheel. DeepEP looks for rocshmem at ROCSHMEM_DIR.
559 +# DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list.
560 +# We pass --rocm-gfx942-fp8-fnuz-max 224 to align the DeepEP dispatch FP8
561 +# quantization bound with vLLM's own clamp
562 +# (vllm/model_executor/layers/quantization/utils/fp8_utils.py): 224 for E4M3FNUZ
563 +RUN --mount=type=cache,target=/root/.cache/ccache \
564 + export PYTORCH_ROCM_ARCH="gfx942;gfx950" \
565 + && git clone ${DEEPEP_REPO} \
566 + && cd DeepEP \
567 + && git checkout ${DEEPEP_BRANCH} \
568 + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} --rocm-gfx942-fp8-fnuz-max 224 bdist_wheel --dist-dir=/app/deep_install
569 +
570 +# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
571 +# not force users to rebuild the long-lived Dockerfile.rocm_base image.
572 +FROM base AS mori_base
573 +ARG NIC_BACKEND
574 +ARG AINIC_VERSION
575 +ARG UBUNTU_CODENAME
576 +RUN /bin/bash -lc 'set -euo pipefail; \
577 + \
578 + install_ainic() { \
579 + apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \
580 + rm -rf /var/lib/apt/lists/*; \
581 + mkdir -p /etc/apt/keyrings; \
582 + curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \
583 + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \
584 + > /etc/apt/sources.list.d/amdainic.list; \
585 + apt-get update && apt-get install -y --no-install-recommends \
586 + libionic-dev \
587 + ionic-common \
588 + ; \
589 + rm -rf /var/lib/apt/lists/*; \
590 + }; \
591 + \
592 + # NOTE: requires FW 235.2.86.0 and kernel drivers on the host: \
593 + # bnxt-en-dkms=1.10.3.235.2.86.0 bnxt-re-dkms=235.2.86.0 (from packages.broadcom.com PPA) \
594 + install_bnxt() { \
595 + install -m 0755 -d /etc/apt/keyrings; \
596 + curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \
597 + -o /etc/apt/keyrings/broadcom-nic.asc; \
598 + chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \
599 + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \
600 + > /etc/apt/sources.list.d/broadcom-nic.list; \
601 + apt-get update && apt-get install -y --no-install-recommends \
602 + bnxt-rocelib=235.2.86.0 \
603 + ; \
604 + cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/; \
605 + ldconfig; \
606 + rm -rf /var/lib/apt/lists/*; \
607 + }; \
608 + \
609 + echo "[MORI] Install MoRI proxy deps"; \
610 + pip install --quiet --ignore-installed blinker && \
611 + pip install --quiet quart msgpack aiohttp pyzmq; \
612 + echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \
613 + \
614 + # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \
615 + # Only vendor packages are installed here for dlopen; no compile-time flags needed. \
616 + case "${NIC_BACKEND}" in \
617 + none) ;; \
618 + all) install_ainic; install_bnxt ;; \
619 + ainic) install_ainic ;; \
620 + bnxt) install_bnxt ;; \
621 + *) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \
622 + esac'
623 +
624 +# -----------------------
625 +# vLLM wheel release build stage (for building distributable wheels)
626 +# This stage pins dependencies to custom ROCm wheel versions and handles version detection
627 +FROM fetch_vllm AS build_vllm_wheel_release
628 +
629 +ARG COMMON_WORKDIR
630 +ARG ROCM_TRITON_KERNELS_COMMIT
631 +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt
632 +COPY --from=rocm-triton-kernels \
633 + /opt/rocm-triton-kernels /opt/rocm-triton-kernels
634 +ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels
635 +
636 +# Drop the pre-built Rust artifacts into the source tree. setup.py detects
637 +# them and ships them as-is, skipping the local Rust build.
638 +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
639 +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
640 +
641 +# Create /install directory for custom wheels
642 +RUN mkdir -p /install
643 +
644 +# Copy custom ROCm wheels from docker/context if they exist
645 +# COPY ensures Docker cache is invalidated when wheels change
646 +# .keep file ensures directory always exists for COPY to work
647 +COPY docker/context/base-wheels/ /tmp/base-wheels/
648 +# This is how we know if we are building for a wheel release or not.
649 +# If there are not wheels found there, we are not building for a wheel release.
650 +# So we exit with an error. To skip this stage.
651 +RUN if [ -n "$(ls /tmp/base-wheels/*.whl 2>/dev/null)" ]; then \
652 + echo "Found custom wheels - copying to /install"; \
653 + cp /tmp/base-wheels/*.whl /install/ && \
654 + echo "Copied custom wheels:"; \
655 + ls -lh /install/; \
656 + else \
657 + echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \
658 + echo "Wheel releases require pre-built ROCm wheels."; \
659 + exit 1; \
660 + fi
661 +
662 +# GIT_REPO_CHECK: Verify repo is clean and tags are available (for release builds)
663 +# This matches CUDA's Dockerfile behavior for proper version detection via setuptools_scm
664 +ARG GIT_REPO_CHECK=0
665 +RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \
666 + echo "Running repository checks..."; \
667 + cd vllm && bash tools/check_repo.sh; \
668 + fi
669 +
670 +# Install the packaging helpers before modifying requirements/rocm.txt.
671 +RUN --mount=type=cache,target=/root/.cache/uv \
672 + grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \
673 + vllm/cmake/external_projects/triton_kernels.cmake \
674 + && cd vllm \
675 + && uv pip install --system setuptools_scm regex
676 +
677 +# Fail if git-based package dependencies are found in requirements files
678 +# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI)
679 +# Extra notes: pip install is able to handle git+ URLs, but uv doesn't.
680 +RUN echo "Checking for git-based packages in requirements files..." \
681 + && echo "Checking common.txt for git-based packages:" \
682 + && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \
683 + echo "ERROR: Git-based packages found in common.txt:"; \
684 + grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \
685 + echo "Please publish these packages to PyPI instead of using git dependencies."; \
686 + exit 1; \
687 + else \
688 + echo " ✓ No git-based packages found in common.txt"; \
689 + fi \
690 + && echo "Checking rocm.txt for git-based packages:" \
691 + && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \
692 + echo "ERROR: Git-based packages found in rocm.txt:"; \
693 + grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \
694 + echo "Please publish these packages to PyPI instead of using git dependencies."; \
695 + exit 1; \
696 + else \
697 + echo " ✓ No git-based packages found in rocm.txt"; \
698 + fi \
699 + && echo "All requirements files are clean - no git-based packages found"
700 +
701 +# Pin vLLM dependencies to exact versions of custom ROCm wheels
702 +# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi
703 +COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py
704 +RUN echo "Pinning vLLM dependencies to custom wheel versions..." \
705 + && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt
706 +
707 +# Install dependencies using custom wheels from /install
708 +RUN --mount=type=cache,target=/root/.cache/uv \
709 + cd vllm \
710 + && echo "Building vLLM with custom wheels from /install" \
711 + && uv pip install --system --find-links /install -r requirements/rocm.txt
712 +
713 +# Build wheel using the shared source version to avoid dirty state from modified requirements/rocm.txt
714 +# (setup.py auto-detects ccache/sccache in PATH)
715 +RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
716 + cd vllm \
717 + && export CCACHE_BASEDIR="$PWD" \
718 + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \
719 + && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \
720 + && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist
721 +
722 +FROM scratch AS export_vllm_wheel_release
723 +ARG COMMON_WORKDIR
724 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/dist/*.whl /
725 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requirements
726 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
727 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests
728 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples
729 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/scripts/benchmark_helion_kernels.py /scripts/benchmark_helion_kernels.py
730 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
731 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile
732 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu
733 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
734 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base
735 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl
736 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl
737 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl
738 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
739 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
740 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust /rust
741 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
742 +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
743 +
744 +# -----------------------
745 +# CI base image (Tier 1) - stable, rarely changing CI dependencies.
746 +# Per-PR test builds pull this as CI_BASE_IMAGE so the test stage only layers
747 +# in the vLLM artifacts for the current commit.
748 +FROM mori_base AS ci_base
749 +ARG COMMON_WORKDIR
750 +
751 +ENV CARGO_HOME=/root/.cargo
752 +ENV RUSTUP_HOME=/root/.rustup
753 +ENV PATH=${CARGO_HOME}/bin:${PATH}
754 +
755 +# Bake the Rust toolchain into ci_base so Rust build/test steps do not
756 +# need to fetch it on each per-commit image build.
757 +COPY --from=rust-toolchain /root/.cargo /root/.cargo
758 +COPY --from=rust-toolchain /root/.rustup /root/.rustup
759 +
760 +# Update rdma-core to support latest rocshmem.
761 +ARG DEEPEP_NIC
762 +RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
763 + git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \
764 + cd /tmp/rdma-core && \
765 + mkdir -p build && cd build && \
766 + cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \
767 + ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
768 +fi
769 +
770 +# Install NIXL + DeepEP wheels.
771 +RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
772 + --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
773 + uv pip install --system /nixl_install/*.whl /deep_install/*.whl
774 +
775 +# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec.
776 +RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
777 + librdmacm1 \
778 + libibverbs1 \
779 + ibverbs-providers \
780 + ibverbs-utils \
781 + util-linux \
782 + unzip \
783 + pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \
784 + libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \
785 + && rm -rf /var/lib/apt/lists/*
786 +
787 +# Install torchcodec from source for ROCm/torch ABI compatibility.
788 +COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh
789 +RUN --mount=type=cache,target=/root/.cache/uv \
790 + --mount=type=cache,target=/root/.cache/pip \
791 + --mount=type=cache,target=/root/.cache/torchcodec-wheels \
792 + bash /tmp/install_torchcodec.sh \
793 + && rm /tmp/install_torchcodec.sh \
794 + && apt-get clean && rm -rf /var/lib/apt/lists/*
795 +
796 +# Pre-install shared ROCm runtime dependencies.
797 +COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/
798 +RUN --mount=type=cache,target=/root/.cache/uv \
799 + uv pip install --system -c /etc/rocm-constraints.txt \
800 + -r /tmp/ci-base-requirements/rocm.txt \
801 + && rm -rf /tmp/ci-base-requirements
802 +
803 +# Enable fast and less brittle model downloads in tests.
804 +ENV HF_XET_HIGH_PERFORMANCE=1
805 +ENV HF_HUB_DOWNLOAD_TIMEOUT=60
806 +
807 +# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py).
808 +ENV PYTORCH_NVML_BASED_CUDA_CHECK=1
809 +
810 +# Pre-install vLLM test dependencies.
811 +COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt
812 +RUN --mount=type=cache,target=/root/.cache/uv \
813 + uv pip install --system -r /tmp/rocm-test-reqs.txt
814 +
815 +# Rebuild fastsafetensors from source so its C++ extension is compiled with
816 +# USE_ROCM and can detect libamdhip64.so at runtime.
817 +RUN --mount=type=cache,target=/root/.cache/pip \
818 + FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \
819 + && test -n "${FASTSAFETENSORS_REQ}" \
820 + && python3 -m pip install --force-reinstall --no-deps \
821 + --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \
822 + && rm /tmp/rocm-test-reqs.txt
823 +
824 +# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel.
825 +# See: https://github.com/pytorch/pytorch/issues/169857
826 +ENV MIOPEN_DEBUG_CONV_DIRECT=0
827 +ENV MIOPEN_DEBUG_CONV_GEMM=0
828 +
829 +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
830 +# See: https://github.com/ROCm/rocm-libraries/issues/6266
831 +ENV HSA_ENABLE_IPC_MODE_LEGACY=1
832 +ENV UCX_RMA_PPLN_ENABLE=y
833 +ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
834 +
835 +# ROCm profiler limits workaround.
836 +RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
837 +ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
838 +
839 +# Install vllm_test_utils in ci_base for ci_base + wheel parity.
840 +COPY tests/vllm_test_utils /tmp/vllm_test_utils
841 +RUN --mount=type=cache,target=/root/.cache/uv \
842 + uv pip install --system /tmp/vllm_test_utils \
843 + && rm -rf /tmp/vllm_test_utils
844 +
845 +RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_install \
846 + --mount=type=bind,source=requirements/test/rocm.txt,target=/tmp/lmcache-constraints.txt \
847 + --mount=type=cache,target=/root/.cache/uv \
848 + uv pip install --system --constraint /tmp/lmcache-constraints.txt \
849 + /lmcache_install/*.whl \
850 + && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)"
851 +
852 +# -----------------------
853 +# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base.
854 +FROM ${CI_BASE_IMAGE} AS test
855 +ARG COMMON_WORKDIR
856 +
857 +# Install the vLLM wheel (--no-deps: all deps already in ci_base).
858 +RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
859 + --mount=type=cache,target=/root/.cache/uv \
860 + cd /install \
861 + && uv pip install --system --no-deps *.whl
862 +
863 +# Store the vLLM wheel in the image for python-only install tests.
864 +COPY --from=export_vllm /*.whl /opt/vllm-wheels/
865 +
866 +WORKDIR /vllm-workspace
867 +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace
868 +
869 +# Copy in the v1 package (for python-only install test group).
870 +COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
871 +
872 +# Hide source under src/ so it won't shadow the installed package in tests.
873 +RUN mkdir src && mv vllm src/vllm
874 +
875 +# Catch GPU<->CPU syncs in execute_model/sample_tokens during tests.
876 +ENV VLLM_GPU_SYNC_CHECK=error
877 +
878 +# -----------------------
879 +# Final vLLM image
880 +FROM mori_base AS final_common
881 +
882 +RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
883 +
884 +# Clean up sccache from release image (not needed at runtime)
885 +# This removes the binary and wrappers that may have been installed during build
886 +RUN rm -f /usr/bin/sccache || true \
887 + && rm -rf /opt/sccache-wrappers || true
888 +
889 +# Unset sccache environment variables for the release image
890 +# This prevents S3 bucket config from leaking into production images
891 +ENV SCCACHE_BUCKET=
892 +ENV SCCACHE_REGION=
893 +ENV SCCACHE_ENDPOINT=
894 +ENV SCCACHE_S3_NO_CREDENTIALS=
895 +ENV SCCACHE_IDLE_TIMEOUT=
896 +
897 +# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt.
898 +# Manually remove it so that later steps of numpy upgrade can continue
899 +RUN case "$(which python3)" in \
900 + *"/opt/conda/envs/py_3.9"*) \
901 + rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \
902 + *) ;; esac
903 +
904 +RUN --mount=type=cache,target=/root/.cache/uv \
905 + uv pip install --system --upgrade huggingface-hub[cli]
906 +
907 +# Install vLLM using uv (inherited from base stage)
908 +# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version
909 +RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
910 + --mount=type=cache,target=/root/.cache/uv \
911 + cd /install \
912 + && uv pip install --system -c /etc/rocm-constraints.txt -r requirements/rocm.txt \
913 + && pip uninstall -y vllm \
914 + && uv pip install --system -c /etc/rocm-constraints.txt *.whl
915 +
916 +# Install NIXL ROCm wheel
917 +RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
918 + uv pip install --system /nixl_install/*.whl
919 +
920 +# Recompute the base's triton fingerprint to detect resolver interference.
921 +RUN EXPECTED=$(grep '^TRITON_SHA256: ' /app/versions.txt | cut -d' ' -f2) \
922 + && ACTUAL=$(find $(python3 -c 'import triton, os; print(os.path.dirname(triton.__file__))') \
923 + -type f \( -name '*.py' -o -name '*.so' \) | sort | xargs sha256sum | sha256sum | cut -d' ' -f1) \
924 + && echo "triton sha256 expected=${EXPECTED} actual=${ACTUAL}" \
925 + && test -n "${EXPECTED}" \
926 + && test "${EXPECTED}" = "${ACTUAL}" \
927 + && python3 -c "import torch; assert torch.version.hip, f'non-ROCm torch: {torch.__version__}'; print(f'torch={torch.__version__}')"
928 +
929 +# Recompute the base's pybind11 fingerprint to detect resolver interference.
930 +RUN EXPECTED=$(grep '^PYBIND11_SHA256: ' /app/versions.txt | cut -d' ' -f2) \
931 + && EXPECTED_ABI=$(grep '^PYBIND11_INTERNALS_VERSION: ' /app/versions.txt | cut -d' ' -f2) \
932 + && INC=$(python3 -c 'import pybind11; print(pybind11.get_include())') \
933 + && ACTUAL=$(find "${INC}" -type f -name '*.h' | sort | xargs sha256sum | sha256sum | cut -d' ' -f1) \
934 + && ACTUAL_ABI=$(grep -rhE '^#[[:space:]]*define PYBIND11_INTERNALS_VERSION' "${INC}" | awk '{print $NF}') \
935 + && echo "pybind11 $(pip show pybind11 | awk '/^Version:/ {print $2}')" \
936 + "sha256 expected=${EXPECTED} actual=${ACTUAL}" \
937 + "internals expected=${EXPECTED_ABI} actual=${ACTUAL_ABI}" \
938 + && test -n "${EXPECTED}" \
939 + && test "${EXPECTED}" = "${ACTUAL}" \
940 + && test "${EXPECTED_ABI}" = "${ACTUAL_ABI}"
941 +
942 +ARG COMMON_WORKDIR
943 +ARG BASE_IMAGE
944 +ARG NIC_BACKEND
945 +ARG AINIC_VERSION
946 +
947 +# Copy over the benchmark scripts as well
948 +COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
949 +COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples
950 +COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
951 +
952 +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
953 +# See: https://github.com/ROCm/rocm-libraries/issues/6266
954 +ENV HSA_ENABLE_IPC_MODE_LEGACY=1
955 +ENV UCX_RMA_PPLN_ENABLE=y
956 +ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
957 +
958 +ENV TOKENIZERS_PARALLELISM=false
959 +
960 +# ENV that can improve safe tensor loading, and end-to-end time
961 +ENV SAFETENSORS_FAST_GPU=1
962 +
963 +# Performance environment variable.
964 +ENV HIP_FORCE_DEV_KERNARG=1
965 +
966 +# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py).
967 +ENV PYTORCH_NVML_BASED_CUDA_CHECK=1
968 +
969 +# Workaround for ROCm profiler limits
970 +RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
971 +ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
972 +RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
973 + && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
974 + && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt
975 +
976 +CMD ["/bin/bash"]
977 +
978 +FROM final_common AS final_lmcache_false
979 +
980 +FROM final_common AS final_lmcache_true
981 +RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_install \
982 + --mount=type=cache,target=/root/.cache/uv \
983 + uv pip install --system /lmcache_install/*.whl \
984 + && test "$(python3 -c 'import numpy; print(numpy.__version__)')" = "2.3.5" \
985 + && uv pip check --system \
986 + && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)"
987 +
988 +FROM final_lmcache_${INSTALL_LMCACHE} AS final
989 +
990 +#Set entrypoint for vllm-openai official images
991 +FROM final AS vllm-openai
992 +ENTRYPOINT ["vllm", "serve"]
docker/Dockerfile.rock_base new
+360
@@ -0,0 +1,360 @@
1 +ARG BASE_IMAGE=ubuntu:22.04
2 +ARG PYTHON_VERSION=3.12
3 +ARG TORCH_VERSION=2.12.0+rocm7.14.0
4 +ARG TORCHVISION_VERSION=0.27.0+rocm7.14.0
5 +ARG TORCHAUDIO_VERSION=2.11.0+rocm7.14.0
6 +ARG TRITON_BRANCH="f0b55c07da61c71775bef6d1a15ebf846430ac75" # release/internal/3.7.x as of 08/13
7 +ARG TRITON_REPO="https://github.com/ROCm/triton.git"
8 +ARG FA_BRANCH="0e60e394"
9 +ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
10 +ARG AITER_BRANCH="v0.1.19"
11 +ARG AITER_REPO="https://github.com/ROCm/aiter.git"
12 +ARG MORI_BRANCH="f95cd5e0ca7b730f92f9ffdfb358484d6965d9f7" # TheRock compatible v1.1.0
13 +ARG MORI_REPO="https://github.com/ROCm/mori.git"
14 +
15 +
16 +# TheRock configuration
17 +ARG SITE_PACKAGES=/usr/local/lib/python${PYTHON_VERSION}/dist-packages
18 +ARG ROCM_RELEASE_WHEELS_MULTIARCH_URL="https://repo.amd.com/rocm/whl-multi-arch/" # For pre-built GPU images, packages, etc.
19 +ARG ROCM_NIGHTLY_WHEELS_MULTIARCH_URL="https://rocm.nightlies.amd.com/whl-multi-arch/" # For torch, triton, etc. wheels with latest fixes.
20 +ARG ROCM_SDK_VERSION=7.14.0
21 +
22 +# Build configuration
23 +ARG PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151"
24 +ARG AITER_ROCM_ARCH="gfx942;gfx950"
25 +
26 +# Sccache configuration (only used in release pipeline)
27 +ARG USE_SCCACHE
28 +ARG SCCACHE_DOWNLOAD_URL
29 +ARG SCCACHE_ENDPOINT
30 +ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
31 +ARG SCCACHE_REGION_NAME=us-west-2
32 +ARG SCCACHE_S3_NO_CREDENTIALS=0
33 +
34 +FROM ${BASE_IMAGE} AS base
35 +
36 +ARG PYTORCH_ROCM_ARCH
37 +ENV PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH}"
38 +ARG AITER_ROCM_ARCH
39 +ENV AITER_ROCM_ARCH="${AITER_ROCM_ARCH}"
40 +
41 +ARG SITE_PACKAGES
42 +ENV SITE_PACKAGES=${SITE_PACKAGES}
43 +ENV SDK_DEV="$SITE_PACKAGES/_rocm_sdk_devel" \
44 + SDK_CORE="$SITE_PACKAGES/_rocm_sdk_core"
45 +ENV ROCM_PATH="$SDK_DEV" \
46 + ROCM_HOME="$SDK_DEV" \
47 + HIP_PATH="$SDK_DEV" \
48 + HIP_CLANG_PATH="$SDK_DEV/lib/llvm/bin" \
49 + HIP_DEVICE_LIB_PATH="$SDK_CORE/lib/llvm/amdgcn/bitcode"
50 +ENV LD_LIBRARY_PATH="$SDK_DEV/lib:$SDK_DEV/lib/rocm_sysdeps/lib:$SDK_CORE/lib:$SDK_CORE/lib/rocm_sysdeps/lib"
51 +ENV PATH="$SDK_DEV/bin:$SDK_DEV/lib/llvm/bin:$SDK_CORE/bin:$PATH"
52 +
53 +# Required for RCCL in ROCm7.1
54 +ENV HSA_NO_SCRATCH_RECLAIM=1
55 +
56 +ARG PYTHON_VERSION
57 +ENV PYTHON_VERSION=${PYTHON_VERSION}
58 +
59 +RUN mkdir -p /app
60 +WORKDIR /app
61 +ENV DEBIAN_FRONTEND=noninteractive
62 +
63 +# Install Python and other dependencies
64 +RUN apt-get update -y \
65 + && apt-get install -y software-properties-common build-essential git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config libdrm-dev libnuma-dev kmod \
66 + && for i in 1 2 3; do \
67 + add-apt-repository -y ppa:deadsnakes/ppa && break || \
68 + { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
69 + done \
70 + && apt-get update -y \
71 + && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \
72 + python${PYTHON_VERSION}-lib2to3 python-is-python3 \
73 + && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
74 + && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
75 + && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
76 + && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \
77 + && python3 --version && python3 -m pip --version
78 +
79 +RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11==3.0.4 Cython
80 +RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/*
81 +
82 +# Install ROCm SDK via TheRock
83 +ARG ROCM_RELEASE_WHEELS_MULTIARCH_URL
84 +ARG ROCM_NIGHTLY_WHEELS_MULTIARCH_URL
85 +ARG ROCM_SDK_VERSION
86 +RUN pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \
87 + --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \
88 + "rocm[libraries,devel]${ROCM_SDK_VERSION:+==${ROCM_SDK_VERSION}}" \
89 + && rocm-sdk init
90 +
91 +# Install PyTorch wheels (nightly for latest torch, release as fallback)
92 +ARG TORCH_VERSION
93 +ARG TORCHAUDIO_VERSION
94 +ARG TORCHVISION_VERSION
95 +RUN pip install \
96 + --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \
97 + --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \
98 + "torch==${TORCH_VERSION}" \
99 + "torchaudio==${TORCHAUDIO_VERSION}" \
100 + "torchvision==${TORCHVISION_VERSION}"
101 +RUN python3 -c "import torch; print(f'torch={torch.__version__}, hip={torch.version.hip}')"
102 +
103 +# Install the AMD-compiled device kernels corresponding to PYTORCH_ROCM_ARCH:
104 +
105 +# Possible values that can be in PYTORCH_ROCM_ARCH:
106 +#
107 +# gfx90a (MI210/250), gfx942 (MI300), gfx950 (MI350), gfx1100 (7900XTX),
108 +# gfx1101 (7800XT), gfx1200 (9070XT), gfx1201 (9070), gfx1150/1151 (Strix APU)
109 +
110 +# For each arch in PYTORCH_ROCM_ARCH, install the precompiled GPU code that
111 +# PyTorch needs to run on that specific chip.
112 +SHELL ["/bin/bash", "-c"]
113 +RUN TORCH_VER=$(python3 -c "import torch; print(torch.__version__)") \
114 + && for arch in $(echo "${PYTORCH_ROCM_ARCH}" | tr ';' ' '); do \
115 + echo "Installing device packages for ${arch}" \
116 + && pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \
117 + --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \
118 + amd-torch-device-${arch}==${TORCH_VER} \
119 + rocm-sdk-device-${arch}==${ROCM_SDK_VERSION}; \
120 + done
121 +
122 +# Remove gfx90a xnack kpacks (ASAN artifacts that cause hipErrorInvalidImage on MI250)
123 +RUN rm -f "${SITE_PACKAGES}/torch/.kpack/torch_gfx90a:xnack+.kpack" \
124 + "${SITE_PACKAGES}/torch/.kpack/torch_gfx90a:xnack-.kpack"
125 +
126 +RUN cd $SITE_PACKAGES/_rocm_sdk_core/share/amd_smi && pip install .
127 +ENV PYTHONPATH=$SITE_PACKAGES/_rocm_sdk_core/share/amd_smi
128 +
129 +# TheRock runtime fixes
130 +RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \
131 + sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \
132 + fi \
133 + && if [ -f "${SITE_PACKAGES}/_rocm_sdk_core/lib/librocm_smi64.so.1" ] \
134 + && [ ! -e "${ROCM_PATH}/lib/librocm_smi64.so.7" ]; then \
135 + ln -sfn "${SITE_PACKAGES}/_rocm_sdk_core/lib/librocm_smi64.so.1" "${ROCM_PATH}/lib/librocm_smi64.so.7"; \
136 + fi \
137 + && printf '%s\n%s\n' "${ROCM_PATH}/lib" "${SITE_PACKAGES}/_rocm_sdk_core/lib" > /etc/ld.so.conf.d/rocm-pip.conf \
138 + && for f in ${SITE_PACKAGES}/_rocm_sdk_core/lib/libamdhip64.so*; do \
139 + bn=$(basename "$f"); \
140 + if [ -f "${ROCM_PATH}/lib/${bn}" ] && [ -f "$f" ] && [ ! -L "$f" ]; then \
141 + rm -f "$f"; \
142 + ln -sfn "${ROCM_PATH}/lib/${bn}" "$f"; \
143 + fi; \
144 + done \
145 + && ldconfig
146 +
147 +# Install sccache if USE_SCCACHE is enabled (for release builds)
148 +ARG USE_SCCACHE
149 +ARG SCCACHE_DOWNLOAD_URL
150 +ARG SCCACHE_ENDPOINT
151 +ARG SCCACHE_BUCKET_NAME
152 +ARG SCCACHE_REGION_NAME
153 +ARG SCCACHE_S3_NO_CREDENTIALS
154 +RUN if [ "$USE_SCCACHE" = "1" ]; then \
155 + echo "Installing sccache..." \
156 + && SCCACHE_ARCH="x86_64" \
157 + && SCCACHE_VERSION="v0.8.1" \
158 + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \
159 + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \
160 + && tar -xzf /tmp/sccache.tar.gz -C /tmp \
161 + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
162 + && chmod +x /usr/bin/sccache \
163 + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \
164 + && sccache --version; \
165 + fi
166 +
167 +# Setup sccache for HIP compilation via HIP_CLANG_PATH
168 +# This creates wrapper scripts in a separate directory and points HIP to use them
169 +# This avoids modifying the original ROCm binaries which can break detection
170 +# NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm)
171 +# Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1
172 +RUN if [ "$USE_SCCACHE" = "1" ]; then \
173 + echo "Setting up sccache wrappers for HIP compilation..." \
174 + && mkdir -p /opt/sccache-wrappers \
175 + && printf '#!/bin/bash\nexec sccache %s/lib/llvm/bin/clang++ "$@"\n' "${ROCM_PATH}" > /opt/sccache-wrappers/clang++ \
176 + && chmod +x /opt/sccache-wrappers/clang++ \
177 + && printf '#!/bin/bash\nexec sccache %s/lib/llvm/bin/clang "$@"\n' "${ROCM_PATH}" > /opt/sccache-wrappers/clang \
178 + && chmod +x /opt/sccache-wrappers/clang \
179 + && echo "sccache wrappers created in /opt/sccache-wrappers"; \
180 + fi
181 +
182 +ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
183 +ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
184 +ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
185 +ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
186 +
187 +
188 +###
189 +### Triton Build
190 +###
191 +FROM base AS build_triton
192 +ARG TRITON_BRANCH
193 +ARG TRITON_REPO
194 +RUN git clone ${TRITON_REPO}
195 +RUN cd triton \
196 + && git checkout ${TRITON_BRANCH} \
197 + && if [ ! -f setup.py ]; then cd python; fi \
198 + && python3 setup.py bdist_wheel --dist-dir=dist \
199 + && mkdir -p /app/install && cp dist/*.whl /app/install
200 +
201 +
202 +###
203 +### MORI Build
204 +###
205 +FROM base AS build_mori
206 +ARG MORI_BRANCH
207 +ARG MORI_REPO
208 +RUN git clone ${MORI_REPO}
209 +RUN cd mori \
210 + && git checkout ${MORI_BRANCH} \
211 + && git submodule update --init --recursive \
212 + && ln -sfn ${ROCM_PATH} /opt/rocm \
213 + && mkdir -p /usr/lib64 && ln -sf /usr/lib/x86_64-linux-gnu/libc.so /usr/lib64/libc.so \
214 + && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl
215 +RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install
216 +
217 +###
218 +### FlashAttention Build
219 +###
220 +FROM base AS build_fa
221 +ARG FA_BRANCH
222 +ARG FA_REPO
223 +ARG USE_SCCACHE
224 +RUN git clone ${FA_REPO}
225 +RUN cd flash-attention \
226 + && git checkout ${FA_BRANCH} \
227 + && git submodule update --init \
228 + && if [ "$USE_SCCACHE" = "1" ]; then \
229 + export HIP_CLANG_PATH=/opt/sccache-wrappers \
230 + && sccache --show-stats; \
231 + fi \
232 + && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \
233 + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
234 +RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install
235 +
236 +
237 +###
238 +### AITER Build
239 +###
240 +FROM base AS build_aiter
241 +ARG AITER_BRANCH
242 +ARG AITER_REPO
243 +ARG USE_SCCACHE
244 +RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO}
245 +RUN cd aiter \
246 + && git submodule update --init --recursive \
247 + && pip install -r requirements.txt
248 +# AITER_USE_SYSTEM_TRITON keeps setup.py off the index triton, which would
249 +# otherwise displace the one we built.
250 +RUN pip install pyyaml && cd aiter \
251 + && if [ "$USE_SCCACHE" = "1" ]; then \
252 + export HIP_CLANG_PATH=/opt/sccache-wrappers \
253 + && sccache --show-stats; \
254 + fi \
255 + && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
256 + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
257 + && ls /app/aiter/dist/*.whl
258 +RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
259 +
260 +
261 +###
262 +### Final Build
263 +###
264 +
265 +# Wheel release stage -
266 +# only includes dependencies used by wheel release pipeline
267 +FROM base AS debs_wheel_release
268 +RUN mkdir /app/debs
269 +RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
270 + cp /install/*.whl /app/debs
271 +RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
272 + cp /install/*.whl /app/debs
273 +RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
274 + cp /install/*.whl /app/debs
275 +
276 +FROM base AS debs
277 +RUN mkdir /app/debs
278 +RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
279 + cp /install/*.whl /app/debs
280 +RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
281 + cp /install/*.whl /app/debs
282 +RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
283 + cp /install/*.whl /app/debs
284 +RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \
285 + cp /install/*.whl /app/debs
286 +
287 +FROM base AS final
288 +
289 +# Remove the torch wheel's Requires-Dist: triton entry in METADATA and the SHA
290 +# hash (its optional) from RECORD. These will cause Python to reinstall
291 +# Triton and a constraints file can't fix it. This doesn't exist for source
292 +# builds, so it wasn't needed there.
293 +RUN TORCH_DIST_INFO=$(ls -d ${SITE_PACKAGES}/torch-*.dist-info) \
294 + && grep -q '^Requires-Dist: triton' "$TORCH_DIST_INFO/METADATA" \
295 + && sed -i '/^Requires-Dist: triton/d' "$TORCH_DIST_INFO/METADATA" \
296 + && sed -i "s|^\(${TORCH_DIST_INFO##*/}/METADATA\),.*|\1,,|" "$TORCH_DIST_INFO/RECORD"
297 +
298 +# Other wheels first, so they resolve against the index triton. Ours goes in last
299 +# with --no-deps, where it can't pull the resolver anywhere.
300 +RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
301 + pip install $(ls /install/*.whl | grep -v '/triton-') \
302 + && pip install --no-deps --force-reinstall /install/triton-*.whl \
303 + # With the pin gone this must be clean, so any complaint is real.
304 + && ! pip check 2>&1 | grep -iE '^(torch|triton) ' \
305 + # Create a constraints file to protect the versions we need, especially
306 + # triton, requirements will try to reinstall repeatedly.
307 + && pip list --format=freeze \
308 + | grep -E '^(torch|triton|torchaudio|torchvision|pybind11)==' > /etc/rocm-constraints.txt \
309 + && echo "ROCm constraints:" && cat /etc/rocm-constraints.txt \
310 + && test "$(wc -l < /etc/rocm-constraints.txt)" -eq 5 \
311 + && grep -q '^pybind11==3\.0\.' /etc/rocm-constraints.txt \
312 + && python3 -c "import torch; assert torch.version.hip, f'non-ROCm torch: {torch.__version__}'; print(f'torch={torch.__version__}')" \
313 + && python3 -c "import triton; print(f'triton={triton.__version__}')" \
314 + && pip list 2>/dev/null | grep -iE '^(torch|triton) '
315 +
316 +ARG BASE_IMAGE
317 +ARG PYTHON_VERSION
318 +ARG ROCM_SDK_VERSION
319 +ARG TORCH_VERSION
320 +ARG TORCHAUDIO_VERSION
321 +ARG TORCHVISION_VERSION
322 +ARG PYTORCH_ROCM_ARCH
323 +ARG AITER_ROCM_ARCH
324 +ARG TRITON_BRANCH
325 +ARG TRITON_REPO
326 +ARG FA_BRANCH
327 +ARG FA_REPO
328 +ARG AITER_BRANCH
329 +ARG AITER_REPO
330 +ARG MORI_BRANCH
331 +ARG MORI_REPO
332 +# Record hashes for Triton and pybind11 so we can verify later.
333 +RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
334 + && echo "PYTHON_VERSION: ${PYTHON_VERSION}" >> /app/versions.txt \
335 + && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \
336 + && echo "TORCH_VERSION: ${TORCH_VERSION}" >> /app/versions.txt \
337 + && echo "TORCHAUDIO_VERSION: ${TORCHAUDIO_VERSION}" >> /app/versions.txt \
338 + && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \
339 + && echo "PYTORCH_ROCM_ARCH: ${PYTORCH_ROCM_ARCH}" >> /app/versions.txt \
340 + && echo "AITER_ROCM_ARCH: ${AITER_ROCM_ARCH}" >> /app/versions.txt \
341 + && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \
342 + && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \
343 + && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \
344 + && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \
345 + && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \
346 + && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \
347 + && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \
348 + && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt \
349 + # Fingerprint our triton so Dockerfile.rocm can prove it survived the build.
350 + && echo "TRITON_VERSION: $(pip show triton | awk '/^Version:/ {print $2}')" >> /app/versions.txt \
351 + && echo "TRITON_SHA256: $(find $(python3 -c 'import triton, os; print(os.path.dirname(triton.__file__))') \
352 + -type f \( -name '*.py' -o -name '*.so' \) | sort | xargs sha256sum | sha256sum | cut -d' ' -f1)" \
353 + >> /app/versions.txt \
354 + && echo "PYBIND11_VERSION: $(pip show pybind11 | awk '/^Version:/ {print $2}')" >> /app/versions.txt \
355 + && echo "PYBIND11_INTERNALS_VERSION: $(grep -rhE '^#[[:space:]]*define PYBIND11_INTERNALS_VERSION' \
356 + $(python3 -c 'import pybind11; print(pybind11.get_include())') | awk '{print $NF}')" \
357 + >> /app/versions.txt \
358 + && echo "PYBIND11_SHA256: $(find $(python3 -c 'import pybind11; print(pybind11.get_include())') \
359 + -type f -name '*.h' | sort | xargs sha256sum | sha256sum | cut -d' ' -f1)" \
360 + >> /app/versions.txt
pyproject.toml
+1
@@ -138,6 +138,7 @@ ignore-hidden = false
138 extend-ignore-identifiers-re = [".*[Uu][Ee][0-9][Mm][0-9].*"]
139
140 [tool.typos.default.extend-identifiers]
141 +a63ede7 = "a63ede7"
142 bbc5b7ede = "bbc5b7ede"
143 NOOPs = "NOOPs"
144 nin_shortcut = "nin_shortcut"
requirements/build/rock.txt new
+19
@@ -0,0 +1,19 @@
1 +# Common dependencies
2 +-r ../common.txt
3 +
4 +
5 +--extra-index-url https://repo.amd.com/rocm/whl-multi-arch/
6 +torch==2.12.0+rocm7.14.0
7 +torchvision==0.27.0+rocm7.14.0
8 +torchaudio==2.11.0+rocm7.14.0
9 +triton==3.7.1+git0263a6a6.rocm7.14.0
10 +cmake>=3.26.1,<4
11 +packaging>=26.2
12 +setuptools>=77.0.3,<80.0.0
13 +setuptools-scm>=8
14 +setuptools-rust>=1.9.0
15 +wheel
16 +jinja2>=3.1.6
17 +amdsmi==7.0.2
18 +timm>=1.0.17
19 +tilelang==0.1.10
requirements/build/rocm.txt
+6 -6
@@ -2,13 +2,13 @@
2 -r ../common.txt
3
4
5 ---extra-index-url https://download.pytorch.org/whl/rocm7.1
6 -torch==2.11.0
7 -torchvision==0.26.0
8 -torchaudio==2.11.0
9 -triton==3.6.0
5 +--extra-index-url https://repo.amd.com/rocm/whl-multi-arch/
6 +torch==2.12.0+rocm7.14.0
7 +torchvision==0.27.0+rocm7.14.0
8 +torchaudio==2.11.0+rocm7.14.0
9 +triton==3.7.1+git0263a6a6
10 cmake>=3.26.1,<4
11 -packaging>=24.2
11 +packaging>=26.2
12 setuptools>=77.0.3,<80.0.0
13 setuptools-scm>=8
14 setuptools-rust>=1.9.0
requirements/test/rocm.in
+1
@@ -39,6 +39,7 @@ transformers==5.15.0
39 tokenizers==0.22.2
40 schemathesis>=4.0.0 # Required for openai schema test
41 # quantization
42 +bitsandbytes==0.50.0
43 buildkite-test-collector==0.1.9
44 tilelang==0.1.10
45
requirements/test/rocm.txt
+5 -1
@@ -72,6 +72,8 @@ azure-storage-blob==12.28.0
72 # via runai-model-streamer-azure
73 backoff==2.2.1
74 # via -r requirements/test/rocm.in
75 +bitsandbytes==0.50.0
76 + # via -r requirements/test/rocm.in
77 black==26.3.1
78 # via datamodel-code-generator
79 blake3==1.0.8
@@ -566,6 +568,7 @@ numpy==2.2.6
568 # -r requirements/test/../common.txt
569 # accelerate
570 # albumentations
571 + # bitsandbytes
572 # bm25s
573 # contourpy
574 # cupy-cuda12x
@@ -691,10 +694,11 @@ outlines-core==0.2.14
694 # via
695 # -c requirements/common.txt
696 # -r requirements/test/../common.txt
694 -packaging==26.0
697 +packaging==26.2
698 # via
699 # -c requirements/rocm.txt
700 # accelerate
701 + # bitsandbytes
702 # black
703 # datasets
704 # evaluate