[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 <Randall.Smith@amd.com>

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 \