[CI] Automate AITER prebuilt upload flow with GitHub actions

.github/scripts/aiter_upload.sh

@@ -0,0 +1,106 @@
+#!/usr/bin/env bash
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+set -euo pipefail
+
+# Inputs for upload (optional):
+#   NVTE_AITER_PREBUILT_BASE_URL - base URL for prebuilts
+#   NVTE_AITER_PREBUILT_UPLOAD_TOKEN - bearer token for Artifactory
+# Optional flag:
+#   --build : build aiter libs before packaging/uploading; default is package-only.
+
+# Derive ROCm version and aiter commit -> cache key
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+ROCM_PATH="${ROCM_PATH:-/opt/rocm}"
+ROCM_VER="$(head -n1 "${ROCM_PATH}/.info/version" | sed -n 's/^\([0-9]\+\.[0-9]\+\).*/\1/p')"
+
+AITER_DIR="${ROOT_DIR}/3rdparty/aiter"
+GIT_CONFIG_GLOBAL="$(mktemp /tmp/gitconfig.XXXXXX)"
+trap 'rm -f "${GIT_CONFIG_GLOBAL}"' EXIT
+git config --global --add safe.directory "${AITER_DIR}" --file "${GIT_CONFIG_GLOBAL}" >/dev/null 2>&1 || true
+AITER_SHA="$(GIT_CONFIG_GLOBAL=${GIT_CONFIG_GLOBAL} git -C "${AITER_DIR}" rev-parse HEAD)"
+
+KEY="rocm-${ROCM_VER}_aiter-${AITER_SHA}"
+CACHE_ROOT="${ROOT_DIR}/build/aiter-prebuilts"
+EXTRACT_DIR="${CACHE_ROOT}/${KEY}"
+OUTPUT_TGZ="/tmp/${KEY}.tar.gz"
+OUTPUT_SHA="/tmp/${KEY}.tar.gz.sha256"
+
+HAS_UPLOAD=0
+if [[ -n "${NVTE_AITER_PREBUILT_BASE_URL:-}" && -n "${NVTE_AITER_PREBUILT_UPLOAD_TOKEN:-}" ]]; then
+  HAS_UPLOAD=1
+fi
+
+# Skip early when remote prebuilt already exists
+REMOTE_URL=""
+if [[ ${HAS_UPLOAD} -eq 1 ]]; then
+  REMOTE_URL="${NVTE_AITER_PREBUILT_BASE_URL}/${KEY}.tar.gz"
+  if curl -sIf "${REMOTE_URL}" >/dev/null; then
+    echo "[aiter-upload] Remote prebuilt already present at ${REMOTE_URL}; nothing to do."
+    exit 0
+  fi
+fi
+
+# Optional build stage (uses GPU_ARCHS if set, else gfx942;gfx950)
+if [[ "${1:-}" == "--build" ]]; then
+  shift
+  ARCHS="${GPU_ARCHS:-gfx942;gfx950}"
+  echo "[AITER-PREBUILT] Building aiter libs for ${ARCHS} ..."
+  rm -rf "${AITER_DIR}/aiter/jit/build"
+  AITER_LOG_MORE=1 \
+  CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=3 \
+  GPU_ARCHS="${ARCHS}" \
+  python3 "${ROOT_DIR}/3rdparty/aiter/op_tests/cpp/mha/compile.py"
+  mkdir -p "${EXTRACT_DIR}"
+  cp "${ROOT_DIR}/3rdparty/aiter/op_tests/cpp/mha/libmha_fwd.so" "${EXTRACT_DIR}/"
+  cp "${ROOT_DIR}/3rdparty/aiter/op_tests/cpp/mha/libmha_bwd.so" "${EXTRACT_DIR}/"
+fi
+
+# Ensure built libs exist
+if [[ ! -f "${EXTRACT_DIR}/libmha_fwd.so" ]]; then
+  echo "[AITER-PREBUILT] Missing libmha_fwd.so in ${EXTRACT_DIR}" >&2
+  exit 1
+fi
+if [[ ! -f "${EXTRACT_DIR}/libmha_bwd.so" ]]; then
+  echo "[AITER-PREBUILT] Missing libmha_bwd.so in ${EXTRACT_DIR}" >&2
+  exit 1
+fi
+
+echo "[AITER-PREBUILT] Packaging ${EXTRACT_DIR} -> ${OUTPUT_TGZ}"
+tar -C "${CACHE_ROOT}" -czf "${OUTPUT_TGZ}" "${KEY}"
+sha256sum "${OUTPUT_TGZ}" | awk '{print $1}' > "${OUTPUT_SHA}"
+
+if [[ ${HAS_UPLOAD} -eq 1 ]]; then
+  echo "[AITER-PREBUILT] Uploading..."
+  COLUMNS=50 curl --progress-bar --fail -X PUT \
+    -H "Authorization: Bearer ${NVTE_AITER_PREBUILT_UPLOAD_TOKEN}" \
+    -T "${OUTPUT_TGZ}" \
+    "${REMOTE_URL}" \
+    -o /dev/null
+  echo "[AITER-PREBUILT] Uploaded tgz to ${REMOTE_URL}"
+
+  # Verify remote SHA256 matches local
+  REMOTE_SHA_TMP="$(mktemp /tmp/aiter_remote_sha.XXXXXX)"
+  trap 'rm -f "${REMOTE_SHA_TMP}"' EXIT
+  if curl -fsSL "${REMOTE_URL}.sha256" -o "${REMOTE_SHA_TMP}"; then
+    REMOTE_SHA_VAL="$(awk '{print $1}' "${REMOTE_SHA_TMP}")"
+    LOCAL_SHA_VAL="$(cat "${OUTPUT_SHA}")"
+    if [[ "${REMOTE_SHA_VAL}" != "${LOCAL_SHA_VAL}" ]]; then
+      echo "[AITER-PREBUILT] Remote SHA256 mismatch!"
+      exit 1
+    else
+      echo "[AITER-PREBUILT] Remote SHA256 verified."
+    fi
+  else
+    echo "[AITER-PREBUILT] Warning: failed to download remote .sha256 for verification." >&2
+  fi
+fi
+
+echo "[AITER-PREBUILT] Artifacts:"
+echo "  tgz: ${OUTPUT_TGZ}"
+echo "  sha: ${OUTPUT_SHA}"
+if [[ ${HAS_UPLOAD} -eq 0 ]]; then
+  echo "[AITER-PREBUILT] To upload, set NVTE_AITER_PREBUILT_BASE_URL and NVTE_AITER_PREBUILT_UPLOAD_TOKEN."
+fi
+

.github/workflows/aiter-prebuilt-upload.yml

@@ -0,0 +1,76 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+name: AITER Prebuilt Upload
+
+on:
+  workflow_dispatch:
+    inputs:
+      gpu_archs:
+        description: "GPU arch list for aiter build"
+        required: true
+        default: "gfx942;gfx950"
+      docker_image:
+        description: "Docker image"
+        required: false
+        default: ""
+
+jobs:
+  upload:
+    runs-on: linux-mi325-8
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.ref }}
+          submodules: recursive
+          fetch-depth: 0
+
+      - name: Resolve docker image
+        id: cfg
+        run: |
+          IMAGE="${{ inputs.docker_image }}"
+          if [ -z "$IMAGE" ]; then
+            IMAGE="${{ vars.DEV_DOCKER_IMAGE }}"
+          fi
+          if [ -z "$IMAGE" ]; then
+            echo "No docker image provided and vars.DEV_DOCKER_IMAGE is empty." >&2
+            exit 1
+          fi
+          echo "image=${IMAGE}" >> $GITHUB_OUTPUT
+
+      - name: Pull docker image
+        run: docker pull ${{ steps.cfg.outputs.image }}
+
+      - name: Run container
+        run: |
+          docker run -dt \
+            --rm \
+            --name te-aiter-upload \
+            --network=host \
+            --device=/dev/dri --device=/dev/kfd \
+            --shm-size=16G \
+            --pid=host \
+            --group-add $(getent group render | cut -d: -f3) \
+            --group-add $(getent group video | cut -d: -f3) \
+            -v "${{ github.workspace }}:/workspace" \
+            -w /workspace \
+            ${{ steps.cfg.outputs.image }}
+
+      - name: Build and upload aiter prebuilt
+        env:
+          NVTE_AITER_PREBUILT_BASE_URL: https://compute-artifactory.amd.com:5000/artifactory/rocm-generic-local/te-ci/aiter-prebuilts
+          NVTE_AITER_PREBUILT_UPLOAD_TOKEN: ${{ secrets.AITER_ARTIFACTORY_TOKEN }}
+        run: |
+          docker exec \
+            -e NVTE_AITER_PREBUILT_BASE_URL=${NVTE_AITER_PREBUILT_BASE_URL} \
+            -e NVTE_AITER_PREBUILT_UPLOAD_TOKEN=${NVTE_AITER_PREBUILT_UPLOAD_TOKEN} \
+            -e GPU_ARCHS="${{ inputs.gpu_archs }}" \
+            te-aiter-upload bash -c 'set -ex
+              if [ -z "${NVTE_AITER_PREBUILT_UPLOAD_TOKEN}" ]; then
+                echo "Missing secrets.AITER_ARTIFACTORY_TOKEN" >&2
+                exit 1
+              fi
+              export HIP_PATH=""
+              .github/scripts/aiter_upload.sh --build
+            '

transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake

@@ -79,46 +79,3 @@ function(download_aiter_prebuilt DOWNLOAD_SUCCESS)
   message(STATUS "[AITER-PREBUILT] Successfully downloaded.")
   set(${DOWNLOAD_SUCCESS} TRUE PARENT_SCOPE)
 endfunction()
-
-# Create prebuilt tgz file to upload
-function(create_upload_files)
-  # Locate .so files
-  if (NOT EXISTS  "${EXTRACT_DIR}/libmha_fwd.so")
-    message(FATAL_ERROR "[AITER-PREBUILT] Missing libmha_fwd.so")
-  endif()
-  if (NOT EXISTS  "${EXTRACT_DIR}/libmha_bwd.so")
-    message(FATAL_ERROR "[AITER-PREBUILT] Missing libmha_bwd.so")
-  endif()
-
-  # Output paths
-  set(OUTPUT_TGZ "/tmp/${KEY}.tar.gz")
-  set(OUTPUT_SHA "/tmp/${KEY}.tar.gz.sha256")
-
-  message(STATUS "[AITER-PREBUILT] Creating prebuilt files...")
-  # Create archive
-  file(ARCHIVE_CREATE
-       OUTPUT "${OUTPUT_TGZ}"
-       PATHS "${KEY}"
-       WORKING_DIRECTORY "${CACHE_ROOT}"
-       FORMAT "gnutar"
-       COMPRESSION "GZip")
-
-  # Compute SHA256
-  file(SHA256 "${OUTPUT_TGZ}" ARCHIVE_HASH)
-  file(WRITE "${OUTPUT_SHA}" "${ARCHIVE_HASH}")
-  message(STATUS "[AITER-PREBUILT] tgz and sha256 files generated successfully:")
-  message(STATUS "  ${OUTPUT_TGZ}")
-  message(STATUS "  ${OUTPUT_SHA}")
-endfunction()
-
-# ------------------------------------------------------
-# Script-mode entry point (to create upload files)
-# Usage: cmake -DACTION=upload -P /path/to/aiter_prebuilt.cmake
-# ------------------------------------------------------
-if (CMAKE_SCRIPT_MODE_FILE)
-  if (DEFINED ACTION AND ACTION STREQUAL "upload")
-    create_upload_files()
-  else()
-    message(FATAL_ERROR "[AITER-PREBUILT] Invalid ACTION=${ACTION}. Use upload.")
-  endif()
-endif()