[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>
rasmith committed
Sep 22, 2026 at 18:05 UTC
111d7d7be2d4cfa439ac0e0cbd2b14a474bfab6e
1 file changed
+37
-7
docker/Dockerfile.rock_base
+37
-7
@@ -3,7 +3,9 @@ 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_VERSION=3.8.0+git4cff872c.rocm10.0.0
6
+ARG TRITON_BRANCH="c0a142ffcbb42f9a67a56bf791cad472be1d6425" # release/internal/3.8.x as of 9/21/2026
7
+ARG TRITON_VERSION=3.8.0+gitc0a142ff
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.22.post1"
@@ -90,19 +92,15 @@ RUN pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \
92
&& rocm-sdk init
93
94
# Install PyTorch wheels (nightly for latest torch, release as fallback)
93
-# Triton is pinned here rather than left to torch's Requires-Dist, so a wheel bump that
94
-# moves the triton torch wants fails resolution instead of quietly changing compilers.
95
ARG TORCH_VERSION
96
ARG TORCHAUDIO_VERSION
97
ARG TORCHVISION_VERSION
98
-ARG TRITON_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}" \
104
- "torchvision==${TORCHVISION_VERSION}" \
105
- "triton==${TRITON_VERSION}"
103
+ "torchvision==${TORCHVISION_VERSION}"
104
105
# Install the AMD-compiled device kernels corresponding to PYTORCH_ROCM_ARCH:
106
@@ -257,6 +255,20 @@ RUN pip install pyyaml && cd aiter \
255
RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
256
257
258
+###
259
+### Triton Build
260
+###
261
+FROM base AS build_triton
262
+ARG TRITON_BRANCH
263
+ARG TRITON_REPO
264
+RUN git clone ${TRITON_REPO}
265
+RUN cd triton \
266
+ && git checkout ${TRITON_BRANCH} \
267
+ && if [ ! -f setup.py ]; then cd python; fi \
268
+ && python3 setup.py bdist_wheel --dist-dir=dist \
269
+ && mkdir -p /app/install && cp dist/*.whl /app/install
270
+
271
+
272
###
273
### ROCr + CLR build
274
###
@@ -320,6 +332,8 @@ RUN mkdir -p /staging/lib \
332
# only includes dependencies used by wheel release pipeline
333
FROM base AS debs_wheel_release
334
RUN mkdir /app/debs
335
+RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
336
+ cp /install/*.whl /app/debs
337
RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
338
cp /install/*.whl /app/debs
339
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 \
341
342
FROM base AS debs
343
RUN mkdir /app/debs
344
+RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
345
+ cp /install/*.whl /app/debs
346
RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
347
cp /install/*.whl /app/debs
348
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 \
351
cp /install/*.whl /app/debs
352
353
FROM base AS final
354
+ARG TRITON_VERSION
355
356
# torch 2.12 implicitly attaches rocprofiler-sdk; its queue interposition hangs at HSA
357
# teardown, so disable it. TODO: drop once ROCm's stock rocprofiler carries the fix.
@@ -359,7 +376,15 @@ RUN set -eux; \
376
377
RUN ldconfig
378
362
-ARG TRITON_VERSION
379
+# This is to prevent resolvers from checking METADATA and RECORD and clobbering
380
+# the Triton source build we installed. Constraints files are unable to help
381
+# with this. However, these entries are optional and can be blanked which
382
+# solves the problem.
383
+RUN TORCH_DIST_INFO=$(ls -d ${SITE_PACKAGES}/torch-*.dist-info) \
384
+ && grep -q '^Requires-Dist: triton' "$TORCH_DIST_INFO/METADATA" \
385
+ && sed -i '/^Requires-Dist: triton/d' "$TORCH_DIST_INFO/METADATA" \
386
+ && sed -i "s|^\(${TORCH_DIST_INFO##*/}/METADATA\),.*|\1,,|" "$TORCH_DIST_INFO/RECORD"
387
+
388
RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
389
pip install /install/*.whl \
390
&& ! pip check 2>&1 | grep -iE '^(torch|triton) ' \
@@ -383,6 +408,8 @@ ARG TORCHAUDIO_VERSION
408
ARG TORCHVISION_VERSION
409
ARG PYTORCH_ROCM_ARCH
410
ARG AITER_ROCM_ARCH
411
+ARG TRITON_BRANCH
412
+ARG TRITON_REPO
413
ARG FA_BRANCH
414
ARG FA_REPO
415
ARG AITER_BRANCH
@@ -391,6 +418,7 @@ ARG MORI_BRANCH
418
ARG MORI_REPO
419
ARG ROCM_SYSTEMS_REPO
420
ARG ROCM_RUNTIME_COMMIT
421
+ARG TRITON_VERSION
422
# Record hashes for Triton and pybind11 so we can verify later.
423
RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
424
&& echo "PYTHON_VERSION: ${PYTHON_VERSION}" >> /app/versions.txt \
@@ -400,6 +428,8 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
428
&& echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \
429
&& echo "PYTORCH_ROCM_ARCH: ${PYTORCH_ROCM_ARCH}" >> /app/versions.txt \
430
&& echo "AITER_ROCM_ARCH: ${AITER_ROCM_ARCH}" >> /app/versions.txt \
431
+ && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \
432
+ && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \
433
&& echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \
434
&& echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \
435
&& echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \