| 1 | ARG BASE_IMAGE=ubuntu:24.04 |
| 2 | ARG ROCM_WHEEL_INDEX=https://rocm.devreleases.amd.com/whl-multi-arch/ |
| 3 | ARG ROCM_SDK_VERSION=7.14.0a20260623 |
| 4 | ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260623 |
| 5 | ARG TORCHVISION_VERSION=0.26.0+rocm7.14.0a20260623 |
| 6 | ARG TORCHAUDIO_VERSION=2.11.0+rocm7.14.0a20260623 |
| 7 | ARG TRITON_VERSION=3.7.1+git110cd8e2.rocm7.14.0a20260623 |
| 8 | ARG APEX_VERSION=1.11.0+rocm7.14.0a20260623 |
| 9 | |
| 10 | |
| 11 | ARG FA_BRANCH="jpvillam/gfx1250_wip" |
| 12 | ARG FA_REPO="https://github.com/jpvillam-amd/flash-attention.git" |
| 13 | ARG AITER_BRANCH="main" |
| 14 | ARG AITER_REPO="https://github.com/ROCm/aiter.git" |
| 15 | ARG MORI_BRANCH="v1.1.0" |
| 16 | ARG MORI_REPO="https://github.com/ROCm/mori.git" |
| 17 | |
| 18 | # Sccache configuration (only used in release pipeline) |
| 19 | ARG USE_SCCACHE |
| 20 | ARG SCCACHE_DOWNLOAD_URL |
| 21 | ARG SCCACHE_ENDPOINT |
| 22 | ARG SCCACHE_BUCKET_NAME=vllm-build-sccache |
| 23 | ARG SCCACHE_REGION_NAME=us-west-2 |
| 24 | ARG SCCACHE_S3_NO_CREDENTIALS=0 |
| 25 | |
| 26 | FROM ${BASE_IMAGE} AS base |
| 27 | |
| 28 | ARG PYTORCH_ROCM_ARCH=gfx1250 |
| 29 | ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} |
| 30 | ENV AITER_ROCM_ARCH=${PYTORCH_ROCM_ARCH} |
| 31 | ENV MORI_GPU_ARCHS=gfx942;gfx950 |
| 32 | ENV FA_GPU_ARCHS=gfx942;gfx950;gfx1250 |
| 33 | |
| 34 | # TODO: Unset these when support is available for gfx1250 |
| 35 | ENV ENABLE_CK=0 |
| 36 | ENV SKIP_CK_BUILD="TRUE" |
| 37 | ENV FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" |
| 38 | ARG PREBUILD_KERNELS=0 |
| 39 | ARG TRITON_KERNEL_BRANCH="padroute" |
| 40 | ARG TRITON_KERNEL_REPO="https://github.com/jpvillam-amd/triton.git" |
| 41 | |
| 42 | # Required for RCCL in ROCm7.1 |
| 43 | ENV HSA_NO_SCRATCH_RECLAIM=1 |
| 44 | |
| 45 | ARG PYTHON_VERSION=3.12 |
| 46 | ENV PYTHON_VERSION=${PYTHON_VERSION} |
| 47 | |
| 48 | RUN mkdir -p /app |
| 49 | WORKDIR /app |
| 50 | ENV DEBIAN_FRONTEND=noninteractive |
| 51 | |
| 52 | # Install Python and other dependencies |
| 53 | RUN apt-get update -y \ |
| 54 | && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev libnuma-dev libdrm-dev pkg-config g++ \ |
| 55 | && for i in 1 2 3; do \ |
| 56 | add-apt-repository -y ppa:deadsnakes/ppa && break || \ |
| 57 | { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ |
| 58 | done \ |
| 59 | && apt-get update -y \ |
| 60 | && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ |
| 61 | python${PYTHON_VERSION}-lib2to3 python-is-python3 \ |
| 62 | && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ |
| 63 | && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ |
| 64 | && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ |
| 65 | && python3 --version |
| 66 | |
| 67 | ENV VIRTUAL_ENV=/opt/venv |
| 68 | RUN python${PYTHON_VERSION} -m venv "${VIRTUAL_ENV}" && \ |
| 69 | "${VIRTUAL_ENV}/bin/python" -m pip install --upgrade pip setuptools PyYAML |
| 70 | ENV PATH=${VIRTUAL_ENV}/bin:$PATH |
| 71 | |
| 72 | RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython |
| 73 | RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* |
| 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 | echo "Installing sccache..." \ |
| 84 | && SCCACHE_ARCH="x86_64" \ |
| 85 | && SCCACHE_VERSION="v0.8.1" \ |
| 86 | && 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}" \ |
| 87 | && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ |
| 88 | && tar -xzf /tmp/sccache.tar.gz -C /tmp \ |
| 89 | && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ |
| 90 | && chmod +x /usr/bin/sccache \ |
| 91 | && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ |
| 92 | && sccache --version; \ |
| 93 | fi |
| 94 | |
| 95 | ## |
| 96 | ## Install PyTorch w/ Triton + ROCM_SDK from ROCM wheel index |
| 97 | ## |
| 98 | ARG ROCM_WHEEL_INDEX |
| 99 | ARG TORCH_VERSION |
| 100 | ARG TORCHVISION_VERSION |
| 101 | ARG TORCHAUDIO_VERSION |
| 102 | ARG ROCM_SDK_VERSION |
| 103 | ARG APEX_VERSION |
| 104 | ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages |
| 105 | ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel |
| 106 | ENV ROCM_HOME=${ROCM_PATH} |
| 107 | ENV ROCM_SOURCE_DIR=${ROCM_PATH} |
| 108 | ENV ROCM_BIN=${ROCM_PATH}/bin |
| 109 | ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake |
| 110 | ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode |
| 111 | ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH |
| 112 | ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib |
| 113 | ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake |
| 114 | ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi |
| 115 | |
| 116 | # torch/torchvision/torchaudio must be pinned to mutually-consistent builds |
| 117 | # (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm |
| 118 | # sdk version is derived from torch's own dependency pin unless overridden, |
| 119 | # which keeps the set consistent and avoids pip backtracking. |
| 120 | RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ |
| 121 | --extra-index-url https://pypi.org/simple \ |
| 122 | "torch[device-all]==${TORCH_VERSION}" \ |
| 123 | "torchvision==${TORCHVISION_VERSION}" \ |
| 124 | "torchaudio==${TORCHAUDIO_VERSION}" \ |
| 125 | "rocm[libraries,devel,device-all]==${ROCM_SDK_VERSION}" && \ |
| 126 | rocm-sdk init |
| 127 | |
| 128 | # Torch runtime deps that may not be published on the ROCm wheel index; |
| 129 | # install them from PyPI afterwards. |
| 130 | RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ |
| 131 | "networkx>=2.5.1" jinja2 "fsspec>=0.8.5" |
| 132 | |
| 133 | |
| 134 | |
| 135 | # Expose the rocm-sdk wheel as a conventional /opt/rocm install so downstream |
| 136 | # builds (Dockerfile.rocm: vLLM csrc, RIXL/UCX, ROCShmem/DeepEP) keep working. |
| 137 | RUN ln -sfn "${ROCM_PATH}" /opt/rocm; |
| 138 | |
| 139 | RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ |
| 140 | sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' \ |
| 141 | "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ |
| 142 | fi |
| 143 | |
| 144 | # The ROCm SDK wheel ships a broken CMake export for hsakmt |
| 145 | # This patch includes the right paths for the numa build target |
| 146 | RUN <<'EOF' |
| 147 | set -eu |
| 148 | TARGETS="/opt/rocm/lib/cmake/hsakmt/hsakmtTargets.cmake" |
| 149 | [ -f "$TARGETS" ] || exit 0 # nothing to patch |
| 150 | grep -q NUMA_LIBRARY "$TARGETS" && exit 0 # already patched |
| 151 | |
| 152 | # 1. Point libdrm's -L at the copy bundled in the wheel, not the builder path. |
| 153 | sed -i 's|-L/__w/[^;"]*|-L${_IMPORT_PREFIX}/lib/rocm_sysdeps/lib|g' "$TARGETS" |
| 154 | |
| 155 | # 2. Drop the nonexistent RHEL libc path (libc is linked implicitly anyway). |
| 156 | sed -i 's|/usr/lib64/libc.so;||g' "$TARGETS" |
| 157 | |
| 158 | # 3. Define the numa::numa target the export references but forgot to create. |
| 159 | cat >> "$TARGETS" <<'CMAKE' |
| 160 | |
| 161 | if(NOT TARGET numa::numa) |
| 162 | find_library(NUMA_LIBRARY NAMES numa REQUIRED) |
| 163 | add_library(numa::numa UNKNOWN IMPORTED) |
| 164 | set_target_properties(numa::numa PROPERTIES IMPORTED_LOCATION "${NUMA_LIBRARY}") |
| 165 | endif() |
| 166 | CMAKE |
| 167 | EOF |
| 168 | |
| 169 | # Clone custom triton_kernels for install later |
| 170 | RUN mkdir -p /app/patched_triton_kernels; \ |
| 171 | cd /app/patched_triton_kernels \ |
| 172 | && git clone ${TRITON_KERNEL_REPO} \ |
| 173 | && cd triton \ |
| 174 | && git checkout ${TRITON_KERNEL_BRANCH} \ |
| 175 | && git submodule update --init --recursive |
| 176 | ENV TRITON_KERNELS_SRC_DIR="/app/patched_triton_kernels/triton/python/triton_kernels/triton_kernels/" |
| 177 | |
| 178 | |
| 179 | # Setup sccache for HIP compilation via HIP_CLANG_PATH |
| 180 | # This creates wrapper scripts in a separate directory and points HIP to use them |
| 181 | # This avoids modifying the original ROCm binaries which can break detection |
| 182 | # NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm) |
| 183 | # Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1 |
| 184 | RUN if [ "$USE_SCCACHE" = "1" ]; then \ |
| 185 | echo "Setting up sccache wrappers for HIP compilation..." \ |
| 186 | && mkdir -p /opt/sccache-wrappers \ |
| 187 | && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ |
| 188 | && chmod +x /opt/sccache-wrappers/clang++ \ |
| 189 | && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ |
| 190 | && chmod +x /opt/sccache-wrappers/clang \ |
| 191 | && echo "sccache wrappers created in /opt/sccache-wrappers"; \ |
| 192 | fi |
| 193 | |
| 194 | # Set sccache environment variables only when USE_SCCACHE=1 |
| 195 | # This prevents S3 config from leaking into images when sccache is not used |
| 196 | ARG USE_SCCACHE |
| 197 | ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} |
| 198 | ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} |
| 199 | ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} |
| 200 | ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} |
| 201 | |
| 202 | |
| 203 | ### |
| 204 | ### AMD SMI Build |
| 205 | ### |
| 206 | FROM base AS build_amdsmi |
| 207 | RUN cd ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi \ |
| 208 | && pip wheel . --wheel-dir=dist |
| 209 | RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/dist/*.whl /app/install |
| 210 | |
| 211 | |
| 212 | ### |
| 213 | ### MORI Build TODO(Build needs fixing) |
| 214 | ### |
| 215 | FROM base AS build_mori |
| 216 | ARG MORI_BRANCH |
| 217 | ARG MORI_REPO |
| 218 | ARG MORI_GPU_ARCHS |
| 219 | RUN mkdir -p /app/install; \ |
| 220 | git clone ${MORI_REPO} \ |
| 221 | && cd mori \ |
| 222 | && git checkout ${MORI_BRANCH} \ |
| 223 | && git submodule update --init --recursive \ |
| 224 | && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl \ |
| 225 | && cp /app/mori/dist/*.whl /app/install; |
| 226 | |
| 227 | |
| 228 | ### |
| 229 | ### FlashAttention Build |
| 230 | ### |
| 231 | # Remove && git submodule update --init \ for CK |
| 232 | FROM base AS build_fa |
| 233 | ARG FA_BRANCH |
| 234 | ARG FA_REPO |
| 235 | ARG USE_SCCACHE |
| 236 | RUN mkdir -p /app/install; \ |
| 237 | git clone ${FA_REPO} \ |
| 238 | && cd flash-attention \ |
| 239 | && git checkout ${FA_BRANCH} \ |
| 240 | && if [ "$USE_SCCACHE" = "1" ]; then \ |
| 241 | export HIP_CLANG_PATH=/opt/sccache-wrappers \ |
| 242 | && sccache --show-stats; \ |
| 243 | fi \ |
| 244 | && GPU_ARCHS=$(echo ${FA_GPU_ARCHS} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ |
| 245 | && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ |
| 246 | && cp dist/*.whl /app/install; |
| 247 | |
| 248 | |
| 249 | ### |
| 250 | ### AITER Build |
| 251 | ### |
| 252 | FROM base AS build_aiter |
| 253 | ARG AITER_BRANCH |
| 254 | ARG AITER_REPO |
| 255 | ARG USE_SCCACHE |
| 256 | RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} |
| 257 | RUN cd aiter \ |
| 258 | && git submodule update --init --recursive \ |
| 259 | && pip install -r requirements.txt |
| 260 | RUN pip install pyyaml && cd aiter \ |
| 261 | && if [ "$USE_SCCACHE" = "1" ]; then \ |
| 262 | export HIP_CLANG_PATH=/opt/sccache-wrappers \ |
| 263 | && sccache --show-stats; \ |
| 264 | fi \ |
| 265 | && AITER_USE_SYSTEM_TRITON=1 PREBUILD_KERNELS=${PREBUILD_KERNELS} GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ |
| 266 | && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ |
| 267 | && ls /app/aiter/dist/*.whl |
| 268 | RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install |
| 269 | |
| 270 | |
| 271 | ### |
| 272 | ### Final Build |
| 273 | ### |
| 274 | |
| 275 | # Wheel release stage - |
| 276 | # only includes dependencies used by wheel release pipeline |
| 277 | FROM base AS debs_wheel_release |
| 278 | RUN mkdir /app/debs |
| 279 | RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ |
| 280 | if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi |
| 281 | RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ |
| 282 | cp /install/*.whl /app/debs |
| 283 | RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ |
| 284 | cp /install/*.whl /app/debs |
| 285 | |
| 286 | # Full debs stage - includes Mori (used by Docker releases) |
| 287 | FROM base AS debs |
| 288 | RUN mkdir /app/debs |
| 289 | RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ |
| 290 | if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi |
| 291 | RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ |
| 292 | cp /install/*.whl /app/debs |
| 293 | RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ |
| 294 | cp /install/*.whl /app/debs |
| 295 | RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ |
| 296 | if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi |
| 297 | |
| 298 | FROM base AS final |
| 299 | RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ |
| 300 | pip install /install/*.whl |
| 301 | |
| 302 | ARG BASE_IMAGE |
| 303 | ARG ROCM_WHEEL_INDEX |
| 304 | ARG ROCM_SDK_VERSION |
| 305 | ARG TORCH_VERSION |
| 306 | ARG TORCHVISION_VERSION |
| 307 | ARG TORCHAUDIO_VERSION |
| 308 | ARG FA_BRANCH |
| 309 | ARG FA_REPO |
| 310 | ARG AITER_BRANCH |
| 311 | ARG AITER_REPO |
| 312 | ARG MORI_BRANCH |
| 313 | ARG MORI_REPO |
| 314 | RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ |
| 315 | && echo "ROCM_WHEEL_INDEX: ${ROCM_WHEEL_INDEX}" >> /app/versions.txt \ |
| 316 | && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \ |
| 317 | && echo "TORCH_VERSION: ${TORCH_VERSION}" >> /app/versions.txt \ |
| 318 | && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \ |
| 319 | && echo "TORCHAUDIO_VERSION: ${TORCHAUDIO_VERSION}" >> /app/versions.txt \ |
| 320 | && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ |
| 321 | && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ |
| 322 | && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ |
| 323 | && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \ |
| 324 | && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \ |
| 325 | && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt |