From 111d7d7be2d4cfa439ac0e0cbd2b14a474bfab6e Mon Sep 17 00:00:00 2001 From: rasmith Date: Tue, 22 Sep 2026 18:05:13 -0500 Subject: [PATCH] [ROCm][Build][The Rock] Bump Triton version to 3.8.x tip-of-tree with source build in The Rock image (#58006) Signed-off-by: Randall Smith --- docker/Dockerfile.rock_base | 44 +++++++++++++++++++++++++++++++------ 1 file changed, 37 insertions(+), 7 deletions(-) diff --git a/docker/Dockerfile.rock_base b/docker/Dockerfile.rock_base index adde16af3c..79896933e7 100644 --- a/docker/Dockerfile.rock_base +++ b/docker/Dockerfile.rock_base @@ -3,7 +3,9 @@ ARG PYTHON_VERSION=3.12 ARG TORCH_VERSION=2.12.0+rocm10.0.0 ARG TORCHVISION_VERSION=0.27.0+rocm10.0.0 ARG TORCHAUDIO_VERSION=2.11.0+rocm10.0.0 -ARG TRITON_VERSION=3.8.0+git4cff872c.rocm10.0.0 +ARG TRITON_BRANCH="c0a142ffcbb42f9a67a56bf791cad472be1d6425" # release/internal/3.8.x as of 9/21/2026 +ARG TRITON_VERSION=3.8.0+gitc0a142ff +ARG TRITON_REPO="https://github.com/ROCm/triton.git" ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" ARG AITER_BRANCH="v0.1.22.post1" @@ -90,19 +92,15 @@ RUN pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ && rocm-sdk init # Install PyTorch wheels (nightly for latest torch, release as fallback) -# Triton is pinned here rather than left to torch's Requires-Dist, so a wheel bump that -# moves the triton torch wants fails resolution instead of quietly changing compilers. ARG TORCH_VERSION ARG TORCHAUDIO_VERSION ARG TORCHVISION_VERSION -ARG TRITON_VERSION RUN pip install \ --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \ "torch==${TORCH_VERSION}" \ "torchaudio==${TORCHAUDIO_VERSION}" \ - "torchvision==${TORCHVISION_VERSION}" \ - "triton==${TRITON_VERSION}" + "torchvision==${TORCHVISION_VERSION}" # Install the AMD-compiled device kernels corresponding to PYTORCH_ROCM_ARCH: @@ -257,6 +255,20 @@ RUN pip install pyyaml && cd aiter \ RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install +### +### Triton Build +### +FROM base AS build_triton +ARG TRITON_BRANCH +ARG TRITON_REPO +RUN git clone ${TRITON_REPO} +RUN cd triton \ + && git checkout ${TRITON_BRANCH} \ + && if [ ! -f setup.py ]; then cd python; fi \ + && python3 setup.py bdist_wheel --dist-dir=dist \ + && mkdir -p /app/install && cp dist/*.whl /app/install + + ### ### ROCr + CLR build ### @@ -320,6 +332,8 @@ RUN mkdir -p /staging/lib \ # only includes dependencies used by wheel release pipeline FROM base AS debs_wheel_release RUN mkdir /app/debs +RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ @@ -327,6 +341,8 @@ RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ FROM base AS debs RUN mkdir /app/debs +RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ @@ -335,6 +351,7 @@ RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ cp /install/*.whl /app/debs FROM base AS final +ARG TRITON_VERSION # torch 2.12 implicitly attaches rocprofiler-sdk; its queue interposition hangs at HSA # teardown, so disable it. TODO: drop once ROCm's stock rocprofiler carries the fix. @@ -359,7 +376,15 @@ RUN set -eux; \ RUN ldconfig -ARG TRITON_VERSION +# This is to prevent resolvers from checking METADATA and RECORD and clobbering +# the Triton source build we installed. Constraints files are unable to help +# with this. However, these entries are optional and can be blanked which +# solves the problem. +RUN TORCH_DIST_INFO=$(ls -d ${SITE_PACKAGES}/torch-*.dist-info) \ + && grep -q '^Requires-Dist: triton' "$TORCH_DIST_INFO/METADATA" \ + && sed -i '/^Requires-Dist: triton/d' "$TORCH_DIST_INFO/METADATA" \ + && sed -i "s|^\(${TORCH_DIST_INFO##*/}/METADATA\),.*|\1,,|" "$TORCH_DIST_INFO/RECORD" + RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ pip install /install/*.whl \ && ! pip check 2>&1 | grep -iE '^(torch|triton) ' \ @@ -383,6 +408,8 @@ ARG TORCHAUDIO_VERSION ARG TORCHVISION_VERSION ARG PYTORCH_ROCM_ARCH ARG AITER_ROCM_ARCH +ARG TRITON_BRANCH +ARG TRITON_REPO ARG FA_BRANCH ARG FA_REPO ARG AITER_BRANCH @@ -391,6 +418,7 @@ ARG MORI_BRANCH ARG MORI_REPO ARG ROCM_SYSTEMS_REPO ARG ROCM_RUNTIME_COMMIT +ARG TRITON_VERSION # Record hashes for Triton and pybind11 so we can verify later. RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ && echo "PYTHON_VERSION: ${PYTHON_VERSION}" >> /app/versions.txt \ @@ -400,6 +428,8 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \ && echo "PYTORCH_ROCM_ARCH: ${PYTORCH_ROCM_ARCH}" >> /app/versions.txt \ && echo "AITER_ROCM_ARCH: ${AITER_ROCM_ARCH}" >> /app/versions.txt \ + && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \ + && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \ && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \