From a631973c768bcdc29cfa841a775192975ba0b241 Mon Sep 17 00:00:00 2001 From: Veera Rajasekhar Reddy Gopu Date: Thu, 15 Jan 2026 06:37:35 +0000 Subject: [PATCH 1/2] CI: add manual aiter prebuilt upload flow Add a workflow_dispatch GHA to build/upload aiter prebuilts using a chosen image and GPU arch list, reusing the shell script. Make ci/aiter_upload.sh handle build + package + upload with optional env-based upload, respecting GPU_ARCHS input and defaulting to gfx942;gfx950. Strip upload/packaging logic out of the CMake helper so normal builds only download/use prebuilts. --- .github/workflows/aiter-prebuilt-upload.yml | 92 +++++++++++++++++++ ci/aiter_upload.sh | 81 ++++++++++++++++ .../common/ck_fused_attn/aiter_prebuilt.cmake | 43 --------- 3 files changed, 173 insertions(+), 43 deletions(-) create mode 100644 .github/workflows/aiter-prebuilt-upload.yml create mode 100755 ci/aiter_upload.sh diff --git a/.github/workflows/aiter-prebuilt-upload.yml b/.github/workflows/aiter-prebuilt-upload.yml new file mode 100644 index 000000000..467140453 --- /dev/null +++ b/.github/workflows/aiter-prebuilt-upload.yml @@ -0,0 +1,92 @@ +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: Host Diagnostics (upload) + run: | + echo "::group::Host Diagnostics" + echo ">>> Active Containers:" + docker ps -a + echo ">>> ROCm Installation:" + ls -d /opt/rocm* || echo "No /opt/rocm found" + echo ">>> GPU info:" + ls -l /dev/dri + ls -l /dev/kfd + rocm-smi || true + echo "::endgroup::" + + - 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 \ + --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: | + if [ -z "${NVTE_AITER_PREBUILT_UPLOAD_TOKEN}" ]; then + echo "Missing secrets.AITER_ARTIFACTORY_TOKEN" >&2 + exit 1 + fi + + 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_INPUT="${{ inputs.gpu_archs }}" \ + te-aiter-upload bash -c 'set -ex + export HIP_PATH="" + export GPU_ARCHS="$GPU_ARCHS_INPUT" + ci/aiter_upload.sh --build + ' + + - name: Cleanup container + if: always() + run: docker rm -f te-aiter-upload || true + diff --git a/ci/aiter_upload.sh b/ci/aiter_upload.sh new file mode 100755 index 000000000..92b5ec281 --- /dev/null +++ b/ci/aiter_upload.sh @@ -0,0 +1,81 @@ +#!/usr/bin/env bash +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 -C "${AITER_DIR}" config --global --add safe.directory "${AITER_DIR}" >/dev/null +AITER_SHA="$(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" + +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 \ + 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}" + +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}" +fi + +echo "[AITER-PREBUILT] Artifacts:" +echo " tgz: ${OUTPUT_TGZ}" +if [[ ${HAS_UPLOAD} -eq 0 ]]; then + echo "[AITER-PREBUILT] To upload, set NVTE_AITER_PREBUILT_BASE_URL and NVTE_AITER_PREBUILT_UPLOAD_TOKEN." +fi + diff --git a/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake b/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake index a4bc9a6be..275e30705 100644 --- a/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake +++ b/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() \ No newline at end of file From 9e8dd63611c2b55dac41262864de55d4e7dc3fc7 Mon Sep 17 00:00:00 2001 From: Veera Rajasekhar Reddy Gopu Date: Sun, 18 Jan 2026 03:58:50 +0000 Subject: [PATCH 2/2] Addressed reviews Move aiter upload helper to .github/scripts, add copyright header, and use a temp gitconfig for safe.directory/commit lookup set CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT, added functionality to verify remote SHA after upload Trim workflow diagnostics/cleanup, use --rm container, pass GPU_ARCHS input directly --- {ci => .github/scripts}/aiter_upload.sh | 31 ++++++++++++++++-- .github/workflows/aiter-prebuilt-upload.yml | 36 ++++++--------------- 2 files changed, 38 insertions(+), 29 deletions(-) rename {ci => .github/scripts}/aiter_upload.sh (67%) diff --git a/ci/aiter_upload.sh b/.github/scripts/aiter_upload.sh similarity index 67% rename from ci/aiter_upload.sh rename to .github/scripts/aiter_upload.sh index 92b5ec281..dd66449b2 100755 --- a/ci/aiter_upload.sh +++ b/.github/scripts/aiter_upload.sh @@ -1,4 +1,7 @@ #!/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): @@ -8,18 +11,21 @@ set -euo pipefail # --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)" +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 -C "${AITER_DIR}" config --global --add safe.directory "${AITER_DIR}" >/dev/null -AITER_SHA="$(git -C "${AITER_DIR}" rev-parse HEAD)" +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 @@ -43,6 +49,7 @@ if [[ "${1:-}" == "--build" ]]; then 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}" @@ -62,6 +69,7 @@ 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..." @@ -71,10 +79,27 @@ if [[ ${HAS_UPLOAD} -eq 1 ]]; then "${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 diff --git a/.github/workflows/aiter-prebuilt-upload.yml b/.github/workflows/aiter-prebuilt-upload.yml index 467140453..d9d9b856c 100644 --- a/.github/workflows/aiter-prebuilt-upload.yml +++ b/.github/workflows/aiter-prebuilt-upload.yml @@ -1,3 +1,6 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. name: AITER Prebuilt Upload on: @@ -23,19 +26,6 @@ jobs: submodules: recursive fetch-depth: 0 - - name: Host Diagnostics (upload) - run: | - echo "::group::Host Diagnostics" - echo ">>> Active Containers:" - docker ps -a - echo ">>> ROCm Installation:" - ls -d /opt/rocm* || echo "No /opt/rocm found" - echo ">>> GPU info:" - ls -l /dev/dri - ls -l /dev/kfd - rocm-smi || true - echo "::endgroup::" - - name: Resolve docker image id: cfg run: | @@ -55,6 +45,7 @@ jobs: - name: Run container run: | docker run -dt \ + --rm \ --name te-aiter-upload \ --network=host \ --device=/dev/dri --device=/dev/kfd \ @@ -71,22 +62,15 @@ jobs: 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: | - if [ -z "${NVTE_AITER_PREBUILT_UPLOAD_TOKEN}" ]; then - echo "Missing secrets.AITER_ARTIFACTORY_TOKEN" >&2 - exit 1 - fi - 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_INPUT="${{ inputs.gpu_archs }}" \ + -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="" - export GPU_ARCHS="$GPU_ARCHS_INPUT" - ci/aiter_upload.sh --build + .github/scripts/aiter_upload.sh --build ' - - - name: Cleanup container - if: always() - run: docker rm -f te-aiter-upload || true -