| 1 | ARG BASE_IMAGE=ubuntu:22.04 |
| 2 | ARG PYTHON_VERSION=3.12 |
| 3 | ARG TORCH_VERSION=2.12.0+rocm10.0.0 |
| 4 | ARG TORCHVISION_VERSION=0.27.0+rocm10.0.0 |
| 5 | ARG TORCHAUDIO_VERSION=2.11.0+rocm10.0.0 |
| 6 | ARG TRITON_BRANCH="669b31ac" # release/internal/3.8.x as of 09/29/2026 |
| 7 | ARG TRITON_VERSION=3.8.0+git669b31ac |
| 8 | ARG TRITON_REPO="https://github.com/ROCm/triton.git" |
| 9 | ARG FA_BRANCH="0e60e394" |
| 10 | ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" |
| 11 | ARG AITER_BRANCH="v0.1.23" |
| 12 | ARG AITER_REPO="https://github.com/ROCm/aiter.git" |
| 13 | ARG MORI_BRANCH="v1.2.3.post1" |
| 14 | ARG MORI_REPO="https://github.com/ROCm/mori.git" |
| 15 | ARG ROCM_SYSTEMS_REPO="https://github.com/ROCm/rocm-systems.git" |
| 16 | ARG ROCM_RUNTIME_COMMIT="25e14349f2606a6b95c4ccd45eb92a1835646cf2" |
| 17 | |
| 18 | |
| 19 | # TheRock configuration |
| 20 | ARG SITE_PACKAGES=/usr/local/lib/python${PYTHON_VERSION}/dist-packages |
| 21 | ARG ROCM_RELEASE_WHEELS_MULTIARCH_URL="https://stable.repo.amd.com/rocm/whl-next/" # For pre-built GPU images, packages, etc. |
| 22 | ARG ROCM_NIGHTLY_WHEELS_MULTIARCH_URL="https://nightly.repo.amd.com/rocm/whl-next/" # For torch, triton, etc. wheels with latest fixes. |
| 23 | ARG ROCM_SDK_VERSION=10.0.0 |
| 24 | |
| 25 | # Build configuration |
| 26 | ARG PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151" |
| 27 | ARG AITER_ROCM_ARCH="gfx942;gfx950" |
| 28 | |
| 29 | # Sccache configuration (only used in release pipeline) |
| 30 | ARG USE_SCCACHE |
| 31 | ARG SCCACHE_DOWNLOAD_URL |
| 32 | ARG SCCACHE_ENDPOINT |
| 33 | ARG SCCACHE_BUCKET_NAME=vllm-build-sccache |
| 34 | ARG SCCACHE_REGION_NAME=us-west-2 |
| 35 | ARG SCCACHE_S3_NO_CREDENTIALS=0 |
| 36 | |
| 37 | FROM ${BASE_IMAGE} AS base |
| 38 | |
| 39 | ARG PYTORCH_ROCM_ARCH |
| 40 | ENV PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH}" |
| 41 | ARG AITER_ROCM_ARCH |
| 42 | ENV AITER_ROCM_ARCH="${AITER_ROCM_ARCH}" |
| 43 | |
| 44 | ARG SITE_PACKAGES |
| 45 | ENV SITE_PACKAGES=${SITE_PACKAGES} |
| 46 | ENV SDK_DEV="$SITE_PACKAGES/_rocm_sdk_devel" \ |
| 47 | SDK_CORE="$SITE_PACKAGES/_rocm_sdk_core" |
| 48 | ENV ROCM_PATH="$SDK_DEV" \ |
| 49 | ROCM_HOME="$SDK_DEV" \ |
| 50 | HIP_PATH="$SDK_DEV" \ |
| 51 | HIP_CLANG_PATH="$SDK_DEV/lib/llvm/bin" \ |
| 52 | HIP_DEVICE_LIB_PATH="$SDK_CORE/lib/llvm/amdgcn/bitcode" |
| 53 | ENV LD_LIBRARY_PATH="$SDK_DEV/lib:$SDK_DEV/lib/rocm_sysdeps/lib:$SDK_CORE/lib:$SDK_CORE/lib/rocm_sysdeps/lib" |
| 54 | ENV PATH="$SDK_DEV/bin:$SDK_DEV/lib/llvm/bin:$SDK_CORE/bin:$PATH" |
| 55 | |
| 56 | # Required for RCCL in ROCm7.1 |
| 57 | ENV HSA_NO_SCRATCH_RECLAIM=1 |
| 58 | |
| 59 | ARG PYTHON_VERSION |
| 60 | ENV PYTHON_VERSION=${PYTHON_VERSION} |
| 61 | |
| 62 | RUN mkdir -p /app |
| 63 | WORKDIR /app |
| 64 | ENV DEBIAN_FRONTEND=noninteractive |
| 65 | |
| 66 | # Install Python and other dependencies |
| 67 | RUN apt-get update -y \ |
| 68 | && 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 \ |
| 69 | && for i in 1 2 3; do \ |
| 70 | add-apt-repository -y ppa:deadsnakes/ppa && break || \ |
| 71 | { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ |
| 72 | done \ |
| 73 | && apt-get update -y \ |
| 74 | && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ |
| 75 | python${PYTHON_VERSION}-lib2to3 python-is-python3 \ |
| 76 | && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ |
| 77 | && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ |
| 78 | && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ |
| 79 | && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ |
| 80 | && python3 --version && python3 -m pip --version |
| 81 | |
| 82 | RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11==3.0.4 Cython |
| 83 | RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* |
| 84 | |
| 85 | # Install ROCm SDK via TheRock |
| 86 | ARG ROCM_RELEASE_WHEELS_MULTIARCH_URL |
| 87 | ARG ROCM_NIGHTLY_WHEELS_MULTIARCH_URL |
| 88 | ARG ROCM_SDK_VERSION |
| 89 | RUN pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ |
| 90 | --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \ |
| 91 | "rocm[libraries,devel]${ROCM_SDK_VERSION:+==${ROCM_SDK_VERSION}}" \ |
| 92 | && rocm-sdk init |
| 93 | |
| 94 | # Install PyTorch wheels (nightly for latest torch, release as fallback) |
| 95 | ARG TORCH_VERSION |
| 96 | ARG TORCHAUDIO_VERSION |
| 97 | ARG TORCHVISION_VERSION |
| 98 | RUN pip install \ |
| 99 | --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ |
| 100 | --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \ |
| 101 | "torch==${TORCH_VERSION}" \ |
| 102 | "torchaudio==${TORCHAUDIO_VERSION}" \ |
| 103 | "torchvision==${TORCHVISION_VERSION}" |
| 104 | |
| 105 | # Install the AMD-compiled device kernels corresponding to PYTORCH_ROCM_ARCH: |
| 106 | |
| 107 | # Possible values that can be in PYTORCH_ROCM_ARCH: |
| 108 | # |
| 109 | # gfx90a (MI210/250), gfx942 (MI300), gfx950 (MI350), gfx1100 (7900XTX), |
| 110 | # gfx1101 (7800XT), gfx1200 (9070XT), gfx1201 (9070), gfx1150/1151 (Strix APU) |
| 111 | |
| 112 | # For each arch in PYTORCH_ROCM_ARCH, install the precompiled GPU code that |
| 113 | # PyTorch needs to run on that specific chip. |
| 114 | SHELL ["/bin/bash", "-c"] |
| 115 | RUN TORCH_VER=$(python3 -c "import torch; print(torch.__version__)") \ |
| 116 | && for arch in $(echo "${PYTORCH_ROCM_ARCH}" | tr ';' ' '); do \ |
| 117 | echo "Installing device packages for ${arch}" \ |
| 118 | && pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ |
| 119 | --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \ |
| 120 | amd-torch-device-${arch}==${TORCH_VER} \ |
| 121 | rocm-sdk-device-${arch}==${ROCM_SDK_VERSION}; \ |
| 122 | done |
| 123 | |
| 124 | # Remove gfx90a xnack kpacks (ASAN artifacts that cause hipErrorInvalidImage on MI250) |
| 125 | RUN rm -f "${SITE_PACKAGES}/torch/.kpack/torch_gfx90a:xnack+.kpack" \ |
| 126 | "${SITE_PACKAGES}/torch/.kpack/torch_gfx90a:xnack-.kpack" |
| 127 | |
| 128 | RUN cd $SITE_PACKAGES/_rocm_sdk_core/share/amd_smi && pip install . |
| 129 | ENV PYTHONPATH=$SITE_PACKAGES/_rocm_sdk_core/share/amd_smi |
| 130 | |
| 131 | # TheRock runtime fixes |
| 132 | RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ |
| 133 | sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ |
| 134 | fi \ |
| 135 | && if [ -f "${SITE_PACKAGES}/_rocm_sdk_core/lib/librocm_smi64.so.1" ] \ |
| 136 | && [ ! -e "${ROCM_PATH}/lib/librocm_smi64.so.7" ]; then \ |
| 137 | ln -sfn "${SITE_PACKAGES}/_rocm_sdk_core/lib/librocm_smi64.so.1" "${ROCM_PATH}/lib/librocm_smi64.so.7"; \ |
| 138 | fi \ |
| 139 | && printf '%s\n%s\n' "${ROCM_PATH}/lib" "${SITE_PACKAGES}/_rocm_sdk_core/lib" > /etc/ld.so.conf.d/rocm-pip.conf \ |
| 140 | && for f in ${SITE_PACKAGES}/_rocm_sdk_core/lib/libamdhip64.so*; do \ |
| 141 | bn=$(basename "$f"); \ |
| 142 | if [ -f "${ROCM_PATH}/lib/${bn}" ] && [ -f "$f" ] && [ ! -L "$f" ]; then \ |
| 143 | rm -f "$f"; \ |
| 144 | ln -sfn "${ROCM_PATH}/lib/${bn}" "$f"; \ |
| 145 | fi; \ |
| 146 | done \ |
| 147 | && for f in librocprofiler-sdk.so.1 librocprofiler-register.so.0 \ |
| 148 | librocprofiler-sdk-roctx.so.1; do \ |
| 149 | if [ -e "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}" ] \ |
| 150 | && [ -e "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" ]; then \ |
| 151 | rm -f "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \ |
| 152 | ln -sfn "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" \ |
| 153 | "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \ |
| 154 | fi; \ |
| 155 | done \ |
| 156 | && ldconfig |
| 157 | |
| 158 | # Install sccache if USE_SCCACHE is enabled (for release builds) |
| 159 | ARG USE_SCCACHE |
| 160 | ARG SCCACHE_DOWNLOAD_URL |
| 161 | ARG SCCACHE_ENDPOINT |
| 162 | ARG SCCACHE_BUCKET_NAME |
| 163 | ARG SCCACHE_REGION_NAME |
| 164 | ARG SCCACHE_S3_NO_CREDENTIALS |
| 165 | RUN if [ "$USE_SCCACHE" = "1" ]; then \ |
| 166 | echo "Installing sccache..." \ |
| 167 | && SCCACHE_ARCH="x86_64" \ |
| 168 | && SCCACHE_VERSION="v0.8.1" \ |
| 169 | && 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}" \ |
| 170 | && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ |
| 171 | && tar -xzf /tmp/sccache.tar.gz -C /tmp \ |
| 172 | && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ |
| 173 | && chmod +x /usr/bin/sccache \ |
| 174 | && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ |
| 175 | && sccache --version; \ |
| 176 | fi |
| 177 | |
| 178 | # Setup sccache for HIP compilation via HIP_CLANG_PATH |
| 179 | # This creates wrapper scripts in a separate directory and points HIP to use them |
| 180 | # This avoids modifying the original ROCm binaries which can break detection |
| 181 | # NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm) |
| 182 | # Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1 |
| 183 | RUN if [ "$USE_SCCACHE" = "1" ]; then \ |
| 184 | echo "Setting up sccache wrappers for HIP compilation..." \ |
| 185 | && mkdir -p /opt/sccache-wrappers \ |
| 186 | && printf '#!/bin/bash\nexec sccache %s/lib/llvm/bin/clang++ "$@"\n' "${ROCM_PATH}" > /opt/sccache-wrappers/clang++ \ |
| 187 | && chmod +x /opt/sccache-wrappers/clang++ \ |
| 188 | && printf '#!/bin/bash\nexec sccache %s/lib/llvm/bin/clang "$@"\n' "${ROCM_PATH}" > /opt/sccache-wrappers/clang \ |
| 189 | && chmod +x /opt/sccache-wrappers/clang \ |
| 190 | && echo "sccache wrappers created in /opt/sccache-wrappers"; \ |
| 191 | fi |
| 192 | |
| 193 | ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} |
| 194 | ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} |
| 195 | ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} |
| 196 | ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} |
| 197 | |
| 198 | |
| 199 | ### |
| 200 | ### MORI Build |
| 201 | ### |
| 202 | FROM base AS build_mori |
| 203 | ARG MORI_BRANCH |
| 204 | ARG MORI_REPO |
| 205 | RUN git clone ${MORI_REPO} |
| 206 | RUN cd mori \ |
| 207 | && git checkout ${MORI_BRANCH} \ |
| 208 | && git submodule update --init --recursive \ |
| 209 | && ln -sfn ${ROCM_PATH} /opt/rocm \ |
| 210 | && pip install "setuptools_scm>=6.2" \ |
| 211 | && CMAKE_PREFIX_PATH="${ROCM_PATH}:${ROCM_PATH}/lib/rocm_sysdeps" \ |
| 212 | python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl |
| 213 | RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install |
| 214 | |
| 215 | ### |
| 216 | ### FlashAttention Build |
| 217 | ### |
| 218 | FROM base AS build_fa |
| 219 | ARG FA_BRANCH |
| 220 | ARG FA_REPO |
| 221 | ARG USE_SCCACHE |
| 222 | RUN git clone ${FA_REPO} |
| 223 | RUN cd flash-attention \ |
| 224 | && git checkout ${FA_BRANCH} \ |
| 225 | && git submodule update --init \ |
| 226 | && if [ "$USE_SCCACHE" = "1" ]; then \ |
| 227 | export HIP_CLANG_PATH=/opt/sccache-wrappers \ |
| 228 | && sccache --show-stats; \ |
| 229 | fi \ |
| 230 | && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ |
| 231 | && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi |
| 232 | RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install |
| 233 | |
| 234 | |
| 235 | ### |
| 236 | ### AITER Build |
| 237 | ### |
| 238 | FROM base AS build_aiter |
| 239 | ARG AITER_BRANCH |
| 240 | ARG AITER_REPO |
| 241 | ARG USE_SCCACHE |
| 242 | RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} |
| 243 | RUN cd aiter \ |
| 244 | && git submodule update --init --recursive \ |
| 245 | && pip install -r requirements.txt |
| 246 | # AITER_USE_SYSTEM_TRITON keeps setup.py from resolving its own triton, which would |
| 247 | # displace the pinned one already installed in base. |
| 248 | RUN pip install pyyaml && cd aiter \ |
| 249 | && if [ "$USE_SCCACHE" = "1" ]; then \ |
| 250 | export HIP_CLANG_PATH=/opt/sccache-wrappers \ |
| 251 | && sccache --show-stats; \ |
| 252 | fi \ |
| 253 | && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ |
| 254 | && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ |
| 255 | && ls /app/aiter/dist/*.whl |
| 256 | RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install |
| 257 | |
| 258 | |
| 259 | ### |
| 260 | ### Triton Build |
| 261 | ### |
| 262 | FROM base AS build_triton |
| 263 | ARG TRITON_BRANCH |
| 264 | ARG TRITON_REPO |
| 265 | RUN git clone ${TRITON_REPO} |
| 266 | RUN cd triton \ |
| 267 | && git checkout ${TRITON_BRANCH} \ |
| 268 | && if [ ! -f setup.py ]; then cd python; fi \ |
| 269 | && python3 setup.py bdist_wheel --dist-dir=dist \ |
| 270 | && mkdir -p /app/install && cp dist/*.whl /app/install |
| 271 | |
| 272 | |
| 273 | ### |
| 274 | ### ROCr + CLR build |
| 275 | ### |
| 276 | # Brings in multi-stream perf improvements; drop after these arrive in an official release |
| 277 | FROM base AS build_rocm_runtime |
| 278 | ARG ROCM_SYSTEMS_REPO |
| 279 | ARG ROCM_RUNTIME_COMMIT |
| 280 | # libgl-dev: rocclr does find_package(OpenGL REQUIRED) unconditionally and the wheel image has |
| 281 | # no GL headers. The link must not pick up a libGL DT_NEEDED, which is asserted after the build. |
| 282 | RUN apt-get update \ |
| 283 | && apt-get install -y --no-install-recommends g++ libdw-dev xxd libgl-dev libglx-dev \ |
| 284 | && rm -rf /var/lib/apt/lists/* |
| 285 | RUN pip install --no-cache-dir CppHeaderParser |
| 286 | RUN git init -q /src && cd /src \ |
| 287 | && git remote add origin ${ROCM_SYSTEMS_REPO} \ |
| 288 | && git config core.sparseCheckout true \ |
| 289 | && git config remote.origin.promisor true \ |
| 290 | && git config remote.origin.partialclonefilter blob:none \ |
| 291 | && git sparse-checkout init --cone \ |
| 292 | && git sparse-checkout set projects/rocr-runtime projects/clr projects/hip shared cmake \ |
| 293 | && git fetch --filter=blob:none --depth 1 origin ${ROCM_RUNTIME_COMMIT} \ |
| 294 | && git checkout -q FETCH_HEAD |
| 295 | # libelf is taken from the SDK's own vendored sysdeps rather than apt -- the wheel image has no |
| 296 | # /usr/include/libelf.h -- so ROCr links against the copy the wheel itself was built against. |
| 297 | ENV SYSDEPS="$SDK_DEV/lib/rocm_sysdeps" |
| 298 | # ROCr first, and installed into $SDK_DEV, because the clr build below needs its headers. |
| 299 | RUN cd /src/projects/rocr-runtime \ |
| 300 | && cmake -B build -G Ninja \ |
| 301 | -DCMAKE_BUILD_TYPE=Release \ |
| 302 | -DCMAKE_PREFIX_PATH="$SDK_DEV;$SYSDEPS" \ |
| 303 | -DCMAKE_INSTALL_PREFIX="$SDK_DEV" \ |
| 304 | -DLIBELF_INCLUDE_DIRS="$SYSDEPS/include" \ |
| 305 | -DLIBELF_LIBRARIES="$SYSDEPS/lib/libelf.so" \ |
| 306 | && cmake --build build --parallel 16 \ |
| 307 | && cmake --install build \ |
| 308 | && cmake --install build --prefix /rocr-install --strip |
| 309 | RUN cd /src/projects/clr \ |
| 310 | && cmake -B build -G Ninja \ |
| 311 | -DCMAKE_BUILD_TYPE=Release \ |
| 312 | -DCLR_BUILD_HIP=ON \ |
| 313 | -DCLR_BUILD_OCL=OFF \ |
| 314 | -DROCM_KPACK_ENABLED=ON \ |
| 315 | -DHIP_COMMON_DIR=/src/projects/hip \ |
| 316 | -DROCM_PATH="$SDK_DEV" \ |
| 317 | -DCMAKE_PREFIX_PATH="$SDK_DEV;$SYSDEPS" \ |
| 318 | -DCMAKE_INSTALL_PREFIX="$SDK_DEV" \ |
| 319 | -DHIP_LLVM_ROOT="$SDK_DEV/lib/llvm" \ |
| 320 | && cmake --build build --parallel 16 \ |
| 321 | && cmake --install build --prefix /clr-install --strip |
| 322 | RUN mkdir -p /staging/lib \ |
| 323 | && cp -P /rocr-install/lib/libhsa-runtime64.so* /staging/lib/ \ |
| 324 | && cp -P /clr-install/lib/libamdhip64.so* /staging/lib/ \ |
| 325 | && ls -l /staging/lib |
| 326 | |
| 327 | |
| 328 | ### |
| 329 | ### Final Build |
| 330 | ### |
| 331 | |
| 332 | # Wheel release stage - |
| 333 | # only includes dependencies used by wheel release pipeline |
| 334 | FROM base AS debs_wheel_release |
| 335 | RUN mkdir /app/debs |
| 336 | RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ |
| 337 | cp /install/*.whl /app/debs |
| 338 | RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ |
| 339 | cp /install/*.whl /app/debs |
| 340 | RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ |
| 341 | cp /install/*.whl /app/debs |
| 342 | |
| 343 | FROM base AS debs |
| 344 | RUN mkdir /app/debs |
| 345 | RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ |
| 346 | cp /install/*.whl /app/debs |
| 347 | RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ |
| 348 | cp /install/*.whl /app/debs |
| 349 | RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ |
| 350 | cp /install/*.whl /app/debs |
| 351 | RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ |
| 352 | cp /install/*.whl /app/debs |
| 353 | |
| 354 | FROM base AS final |
| 355 | ARG TRITON_VERSION |
| 356 | |
| 357 | # torch 2.12 implicitly attaches rocprofiler-sdk; its queue interposition hangs at HSA |
| 358 | # teardown, so disable it. TODO: drop once ROCm's stock rocprofiler carries the fix. |
| 359 | ENV ROCPROFILER_QUEUE_INTERPOSITION=0 |
| 360 | |
| 361 | COPY --from=build_rocm_runtime /staging/lib/ ${SDK_DEV}/lib/ |
| 362 | # Hard link the patched library onto every stock file name, in both SDK trees, so that the |
| 363 | # SONAME lookup and the hard coded paths all resolve to it. Hard links rather than symlinks |
| 364 | # because ldconfig rewrites SONAME symlinks but cannot re-point an inode. $SDK_CORE keeps its |
| 365 | # own libhsa-runtime64 copy, which is on LD_LIBRARY_PATH and in ld.so.conf; relinking it is |
| 366 | # belt and braces rather than strictly required, but it leaves no stock ROCr in the image. |
| 367 | RUN set -eux; \ |
| 368 | cd "$SDK_DEV/lib"; \ |
| 369 | ours_hsa="$(readlink -f libhsa-runtime64.so.1)"; \ |
| 370 | ours_hip="$(readlink -f libamdhip64.so.7)"; \ |
| 371 | for f in "$SDK_DEV"/lib/libhsa-runtime64.so.1* "$SDK_CORE"/lib/libhsa-runtime64.so.1* \ |
| 372 | "$SDK_DEV"/lib/libamdhip64.so.7* "$SDK_CORE"/lib/libamdhip64.so.7*; do \ |
| 373 | [ -f "$f" ] && [ ! -L "$f" ] || continue; \ |
| 374 | case "${f##*/}" in libhsa*) ours="$ours_hsa";; *) ours="$ours_hip";; esac; \ |
| 375 | [ "$f" -ef "$ours" ] || ln -f "$ours" "$f"; \ |
| 376 | done |
| 377 | |
| 378 | RUN ldconfig |
| 379 | |
| 380 | # This is to prevent resolvers from checking METADATA and RECORD and clobbering |
| 381 | # the Triton source build we installed. Constraints files are unable to help |
| 382 | # with this. However, these entries are optional and can be blanked which |
| 383 | # solves the problem. |
| 384 | RUN TORCH_DIST_INFO=$(ls -d ${SITE_PACKAGES}/torch-*.dist-info) \ |
| 385 | && grep -q '^Requires-Dist: triton' "$TORCH_DIST_INFO/METADATA" \ |
| 386 | && sed -i '/^Requires-Dist: triton/d' "$TORCH_DIST_INFO/METADATA" \ |
| 387 | && sed -i "s|^\(${TORCH_DIST_INFO##*/}/METADATA\),.*|\1,,|" "$TORCH_DIST_INFO/RECORD" |
| 388 | |
| 389 | RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ |
| 390 | pip install /install/*.whl \ |
| 391 | && ! pip check 2>&1 | grep -iE '^(torch|triton) ' \ |
| 392 | && test "$(pip show triton | awk '/^Version:/ {print $2}')" = "${TRITON_VERSION}" \ |
| 393 | # Create a constraints file to protect the versions we need, especially |
| 394 | # triton, requirements will try to reinstall repeatedly. |
| 395 | && pip list --format=freeze \ |
| 396 | | grep -E '^(torch|triton|torchaudio|torchvision|pybind11)==' > /etc/rocm-constraints.txt \ |
| 397 | && echo "ROCm constraints:" && cat /etc/rocm-constraints.txt \ |
| 398 | && test "$(wc -l < /etc/rocm-constraints.txt)" -eq 5 \ |
| 399 | && grep -q '^pybind11==3\.0\.' /etc/rocm-constraints.txt \ |
| 400 | && python3 -c "import torch; assert torch.version.hip, f'non-ROCm torch: {torch.__version__}'; print(f'torch={torch.__version__}')" \ |
| 401 | && python3 -c "import triton; print(f'triton={triton.__version__}')" \ |
| 402 | && pip list 2>/dev/null | grep -iE '^(torch|triton) ' |
| 403 | |
| 404 | ARG BASE_IMAGE |
| 405 | ARG PYTHON_VERSION |
| 406 | ARG ROCM_SDK_VERSION |
| 407 | ARG TORCH_VERSION |
| 408 | ARG TORCHAUDIO_VERSION |
| 409 | ARG TORCHVISION_VERSION |
| 410 | ARG PYTORCH_ROCM_ARCH |
| 411 | ARG AITER_ROCM_ARCH |
| 412 | ARG TRITON_BRANCH |
| 413 | ARG TRITON_REPO |
| 414 | ARG FA_BRANCH |
| 415 | ARG FA_REPO |
| 416 | ARG AITER_BRANCH |
| 417 | ARG AITER_REPO |
| 418 | ARG MORI_BRANCH |
| 419 | ARG MORI_REPO |
| 420 | ARG ROCM_SYSTEMS_REPO |
| 421 | ARG ROCM_RUNTIME_COMMIT |
| 422 | ARG TRITON_VERSION |
| 423 | # Record hashes for Triton and pybind11 so we can verify later. |
| 424 | RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ |
| 425 | && echo "PYTHON_VERSION: ${PYTHON_VERSION}" >> /app/versions.txt \ |
| 426 | && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \ |
| 427 | && echo "TORCH_VERSION: ${TORCH_VERSION}" >> /app/versions.txt \ |
| 428 | && echo "TORCHAUDIO_VERSION: ${TORCHAUDIO_VERSION}" >> /app/versions.txt \ |
| 429 | && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \ |
| 430 | && echo "PYTORCH_ROCM_ARCH: ${PYTORCH_ROCM_ARCH}" >> /app/versions.txt \ |
| 431 | && echo "AITER_ROCM_ARCH: ${AITER_ROCM_ARCH}" >> /app/versions.txt \ |
| 432 | && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \ |
| 433 | && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \ |
| 434 | && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ |
| 435 | && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ |
| 436 | && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ |
| 437 | && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \ |
| 438 | && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \ |
| 439 | && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt \ |
| 440 | && echo "ROCM_SYSTEMS_REPO: ${ROCM_SYSTEMS_REPO}" >> /app/versions.txt \ |
| 441 | && echo "ROCM_RUNTIME_COMMIT: ${ROCM_RUNTIME_COMMIT}" >> /app/versions.txt \ |
| 442 | # Fingerprint our triton so Dockerfile.rocm can prove it survived the build. |
| 443 | && echo "TRITON_VERSION: $(pip show triton | awk '/^Version:/ {print $2}')" >> /app/versions.txt \ |
| 444 | && echo "TRITON_SHA256: $(find $(python3 -c 'import triton, os; print(os.path.dirname(triton.__file__))') \ |
| 445 | -type f \( -name '*.py' -o -name '*.so' \) | sort | xargs sha256sum | sha256sum | cut -d' ' -f1)" \ |
| 446 | >> /app/versions.txt \ |
| 447 | && echo "PYBIND11_VERSION: $(pip show pybind11 | awk '/^Version:/ {print $2}')" >> /app/versions.txt \ |
| 448 | && echo "PYBIND11_INTERNALS_VERSION: $(grep -rhE '^#[[:space:]]*define PYBIND11_INTERNALS_VERSION' \ |
| 449 | $(python3 -c 'import pybind11; print(pybind11.get_include())') | awk '{print $NF}')" \ |
| 450 | >> /app/versions.txt \ |
| 451 | && echo "PYBIND11_SHA256: $(find $(python3 -c 'import pybind11; print(pybind11.get_include())') \ |
| 452 | -type f -name '*.h' | sort | xargs sha256sum | sha256sum | cut -d' ' -f1)" \ |
| 453 | >> /app/versions.txt |