From ac593fed901263911bb9cf7564d9e09949ed0345 Mon Sep 17 00:00:00 2001
From: Duyi-Wang <duyi.wang@amd.com>
Date: Fri, 3 Apr 2026 16:54:28 +0800
Subject: [PATCH] [AMD][Dockerfile] Support build-arg AITER_COMMIT for
 rocm.Dockerfile (#21949)

---
 docker/rocm.Dockerfile                      | 10 ++++++----
 scripts/ci/amd/amd_ci_install_dependency.sh |  8 ++++----
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile
index d53f657f0..e8d399f41 100644
--- a/docker/rocm.Dockerfile
+++ b/docker/rocm.Dockerfile
@@ -27,7 +27,7 @@ ENV BUILD_TRITON="0"
 ENV BUILD_LLVM="0"
 ENV BUILD_AITER_ALL="1"
 ENV BUILD_MOONCAKE="1"
-ENV AITER_COMMIT="v0.1.11.post1"
+ENV AITER_COMMIT_DEFAULT="v0.1.11.post1"
 
 # ===============================
 # Base image 942 with rocm720 and args
@@ -37,7 +37,7 @@ ENV BUILD_TRITON="1"
 ENV BUILD_LLVM="0"
 ENV BUILD_AITER_ALL="1"
 ENV BUILD_MOONCAKE="1"
-ENV AITER_COMMIT="v0.1.11.post1"
+ENV AITER_COMMIT_DEFAULT="v0.1.11.post1"
 
 # ===============================
 # Base image 950 and args
@@ -47,7 +47,7 @@ ENV BUILD_TRITON="0"
 ENV BUILD_LLVM="0"
 ENV BUILD_AITER_ALL="1"
 ENV BUILD_MOONCAKE="1"
-ENV AITER_COMMIT="v0.1.11.post1"
+ENV AITER_COMMIT_DEFAULT="v0.1.11.post1"
 
 # ===============================
 # Base image 950 with rocm720 and args
@@ -57,7 +57,7 @@ ENV BUILD_TRITON="1"
 ENV BUILD_LLVM="0"
 ENV BUILD_AITER_ALL="1"
 ENV BUILD_MOONCAKE="1"
-ENV AITER_COMMIT="v0.1.11.post1"
+ENV AITER_COMMIT_DEFAULT="v0.1.11.post1"
 
 # ===============================
 # Chosen arch and args
@@ -79,6 +79,8 @@ ARG TRITON_REPO="https://github.com/triton-lang/triton.git"
 ARG TRITON_COMMIT="42270451990532c67e69d753fbd026f28fcc4840"
 
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
+ARG AITER_COMMIT=""
+ENV AITER_COMMIT="${AITER_COMMIT:-${AITER_COMMIT_DEFAULT}}"
 
 ARG LLVM_REPO="https://github.com/jrbyrnes/llvm-project.git"
 ARG LLVM_BRANCH="MainOpSelV2"
diff --git a/scripts/ci/amd/amd_ci_install_dependency.sh b/scripts/ci/amd/amd_ci_install_dependency.sh
index 9d870d1dd..e3a96f564 100755
--- a/scripts/ci/amd/amd_ci_install_dependency.sh
+++ b/scripts/ci/amd/amd_ci_install_dependency.sh
@@ -200,15 +200,15 @@ echo "[CI-AITER-CHECK] Runner GPU_ARCH=${GPU_ARCH}"
 if [[ "${GPU_ARCH}" == "mi35x" ]]; then
     echo "[CI-AITER-CHECK] Using gfx950 block from Dockerfile..."
     REPO_AITER_COMMIT=$(grep -F -A20 'FROM $BASE_IMAGE_950 AS gfx950' docker/rocm.Dockerfile \
-                        | grep 'AITER_COMMIT=' \
+                        | grep 'AITER_COMMIT_DEFAULT=' \
                         | head -n1 \
-                        | sed 's/.*AITER_COMMIT="\([^"]*\)".*/\1/')
+                        | sed 's/.*AITER_COMMIT_DEFAULT="\([^"]*\)".*/\1/')
 else
     echo "[CI-AITER-CHECK] Using gfx942 block from Dockerfile..."
     REPO_AITER_COMMIT=$(grep -F -A20 'FROM $BASE_IMAGE_942 AS gfx942' docker/rocm.Dockerfile \
-                        | grep 'AITER_COMMIT=' \
+                        | grep 'AITER_COMMIT_DEFAULT=' \
                         | head -n1 \
-                        | sed 's/.*AITER_COMMIT="\([^"]*\)".*/\1/')
+                        | sed 's/.*AITER_COMMIT_DEFAULT="\([^"]*\)".*/\1/')
 fi