From b8dc150ffeff7c47eff71f6624e0679163b2a939 Mon Sep 17 00:00:00 2001
From: Yi DING <yi.ding@amd.com>
Date: Mon, 2 Mar 2026 23:35:22 +0800
Subject: [PATCH] [CK] Fix aiter tests in CI (#5011)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Motivation

Updates the CK/AITER CI Docker build to source Composable Kernel either
from `ROCm/rocm-libraries` (via sparse-checkout) or directly from
`ROCm/composable_kernel`, aiming to make aiter tests reliable in CI.

**Changes:**
- Added a build arg to toggle fetching CK from `ROCm/rocm-libraries`
(enabled by default).
- Implemented sparse-checkout + local re-init/commit flow to materialize
CK into a local `ck/` directory.
- Updated aiter’s CK vendoring step to clone from the locally prepared
`ck/` directory.



## Technical Details

<!-- Explain the changes along with any relevant GitHub links. -->

## Test Plan

<!-- Explain any relevant testing done to verify this PR. -->

## Test Result

<!-- Briefly summarize test outcomes. -->

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
---
 Dockerfile.aiter | 27 +++++++++++++++++++++++----
 1 file changed, 23 insertions(+), 4 deletions(-)

diff --git a/Dockerfile.aiter b/Dockerfile.aiter
index 020afeccf4..a5a3f81fca 100644
--- a/Dockerfile.aiter
+++ b/Dockerfile.aiter
@@ -2,16 +2,35 @@ ARG BASE_DOCKER="rocm/pytorch:latest"
 FROM $BASE_DOCKER
 ARG AITER_BRANCH="main"
 ARG CK_AITER_BRANCH="develop"
+# CK_FROM_ROCM_LIBRARIES - 1: CK from rocm-libraries sparse-checkout; 0: direct clone from ROCm/composable_kernel
+ARG CK_FROM_ROCM_LIBRARIES=1
 RUN pip install pandas zmq einops ninja tabulate && \
     pip install numpy==1.26.2 && \
     sudo mkdir /home/jenkins && \
     sudo mkdir /home/jenkins/workspace && \
-    cd /home/jenkins/workspace && \
-    rm -rf aiter && \
-    git clone -b "$AITER_BRANCH" --recursive https://github.com/ROCm/aiter.git && \
+    cd /home/jenkins/workspace && rm -rf rocm-libraries ck && \
+    if [ "$CK_FROM_ROCM_LIBRARIES" = "1" ]; then \
+        git clone --depth 1 -b "$CK_AITER_BRANCH" --no-checkout --filter=blob:none https://github.com/ROCm/rocm-libraries.git && \
+        cd rocm-libraries && \
+        git sparse-checkout init --cone && \
+        git sparse-checkout set projects/composablekernel && \
+        git checkout "$CK_AITER_BRANCH" && \
+        ROCM_LIBRARIES_SHA=$(git rev-parse --short HEAD) && \
+        mv projects/composablekernel ../ck && \
+        cd ../ck && rm -rf ../rocm-libraries && \
+        git init && \
+        git config user.name "assistant-librarian[bot]" && \
+        git config user.email "assistant-librarian[bot]@users.noreply.github.com" && \
+        git branch -m "$CK_AITER_BRANCH" && git add -A && \
+        git commit -m "import from ROCm/rocm-libraries@$ROCM_LIBRARIES_SHA" ; \
+    else \
+        git clone --depth 1 -b "$CK_AITER_BRANCH" https://github.com/ROCm/composable_kernel.git ck ; \
+    fi && \
+    cd /home/jenkins/workspace && rm -rf aiter && \
+    git clone --depth 1 -b "$AITER_BRANCH" --recursive https://github.com/ROCm/aiter.git && \
     cd aiter && \
     rm -rf 3rdparty/composable_kernel/ && \
-    git clone -b "$CK_AITER_BRANCH" https://github.com/ROCm/composable_kernel.git 3rdparty/composable_kernel/ && \
+    git clone -b "$CK_AITER_BRANCH" ../ck 3rdparty/composable_kernel/ && \
     python3 setup.py develop && \
     groupadd -g 1001 jenkins && \
     useradd -u 1001 -g 1001 -m -s /bin/bash jenkins && \