From c29c9f4e0205d20b82cc3504535c1d82f050abeb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniele=20Trifir=C3=B2?= <dtrifiro@redhat.com>
Date: Wed, 25 Sep 2024 16:55:11 +0200
Subject: [PATCH 1/3] Dockerfile.rocm.ubi: move microdnf update in base stage

---
 Dockerfile.rocm.ubi | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/Dockerfile.rocm.ubi b/Dockerfile.rocm.ubi
index 17909b97f50b2..39093f4630e5c 100644
--- a/Dockerfile.rocm.ubi
+++ b/Dockerfile.rocm.ubi
@@ -13,6 +13,7 @@ ENV VIRTUAL_ENV=/opt/vllm
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"
 
 RUN --mount=type=cache,target=/root/.cache/pip \
+ microdnf -y update && \
  microdnf install -y --setopt=install_weak_deps=0 --nodocs \
     python${PYTHON_VERSION}-devel \
     python${PYTHON_VERSION}-pip \
@@ -40,8 +41,7 @@ gpgcheck=1\n\
 gpgkey=https://repo.radeon.com/rocm/rocm.gpg.key" > /etc/yum.repos.d/amdgpu.repo
 
 
-RUN microdnf -y update && \
-    microdnf -y install \
+RUN microdnf -y install \
         rocm-hip-libraries rocm-hip-runtime \
         miopen-hip && \
     microdnf clean all

From 69ac6c1e66f9f032059cbb732cef4bf088142eb1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniele=20Trifir=C3=B2?= <dtrifiro@redhat.com>
Date: Wed, 25 Sep 2024 16:55:34 +0200
Subject: [PATCH 2/3] Dockerfile.rocm.ubi: bump torch version to
 2.5.0.dev20240912+rocm6.1

---
 Dockerfile.rocm.ubi | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/Dockerfile.rocm.ubi b/Dockerfile.rocm.ubi
index 39093f4630e5c..d6383a30eb40e 100644
--- a/Dockerfile.rocm.ubi
+++ b/Dockerfile.rocm.ubi
@@ -49,9 +49,8 @@ RUN microdnf -y install \
 RUN --mount=type=cache,target=/root/.cache/pip \
     --mount=type=cache,target=/root/.cache/uv \
     uv pip install -v --index-url "https://download.pytorch.org/whl/nightly/rocm6.1" \
-        torch==2.5.0.dev20240726+rocm6.1 \
-        torchvision==0.20.0.dev20240726+rocm6.1
-
+        torch==2.5.0.dev20240912+rocm6.1 \
+        torchvision==0.20.0.dev20240912+rocm6.1
 
 FROM rocm_base as rocm_devel
 

From 399c11499e666be94a954ece49323f644860b2e9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniele=20Trifir=C3=B2?= <dtrifiro@redhat.com>
Date: Wed, 25 Sep 2024 16:56:13 +0200
Subject: [PATCH 3/3] Dockerfile.rocm.ubi: get rid of build triton stage

this is a torch dependency when installed from the pytorch/rocm6.1
index: https://download.pytorch.org/whl/nightly/rocm6.1
---
 Dockerfile.rocm.ubi | 27 ---------------------------
 1 file changed, 27 deletions(-)

diff --git a/Dockerfile.rocm.ubi b/Dockerfile.rocm.ubi
index d6383a30eb40e..12a1da4ca3181 100644
--- a/Dockerfile.rocm.ubi
+++ b/Dockerfile.rocm.ubi
@@ -126,31 +126,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 
 ##################################################################################################
 
-FROM rocm_devel AS build_triton
-
-# Whether to build triton on rocm
-ARG BUILD_TRITON="0"
-ARG TRITON_BRANCH="e0fc12c"
-
-# Build triton wheel if `BUILD_TRITON` is set to 1
-RUN --mount=type=cache,target=/root/.cache/ccache \
-    --mount=type=cache,target=/root/.cache/pip \
-    --mount=type=cache,target=/root/.cache/uv \
-    if [ "$BUILD_TRITON" = "1" ]; then \
-        mkdir -p /libs && cd /libs && \
-        git clone https://github.com/OpenAI/triton.git && \
-        cd triton && \
-        git checkout "${TRITON_BRANCH}" && \
-        cd python && \
-        uv pip install -v cmake ninja && \
-        python3 setup.py bdist_wheel --dist-dir=/install; \
-    else \
-        mkdir -p /install && \
-        wget -P /install "https://download.pytorch.org/whl/pytorch_triton_rocm-3.0.0-cp312-cp312-linux_x86_64.whl#sha256=650ce7a6624ecbbe90c30fccb472a57772c08bc6f12312246f7af584f275b7ca"; \
-    fi
-
-##################################################################################################
-
 FROM rocm_devel AS build_vllm
 ARG PYTORCH_ROCM_ARCH
 ARG MAX_JOBS
@@ -188,7 +163,6 @@ RUN microdnf install -y --setopt=install_weak_deps=0 --nodocs gcc && \
     microdnf clean all
 
 RUN --mount=type=bind,from=build_amdsmi,src=/install,target=/install/amdsmi/ \
-    --mount=type=bind,from=build_triton,src=/install,target=/install/triton \
     --mount=type=bind,from=build_flashattention,src=/install,target=/install/flashattention \
     --mount=type=bind,from=build_vllm,src=/workspace/dist,target=/install/vllm/ \
     --mount=type=cache,target=/root/.cache/pip \
@@ -197,7 +171,6 @@ RUN --mount=type=bind,from=build_amdsmi,src=/install,target=/install/amdsmi/ \
         --index-strategy=unsafe-best-match \
         --extra-index-url "https://download.pytorch.org/whl/nightly/rocm6.1" \
         /install/amdsmi/*.whl\
-        /install/triton/*.whl\
         /install/flashattention/*.whl\
         /install/vllm/*.whl