diff --git a/.github/scripts/aiter_upload.sh b/.github/scripts/aiter_upload.sh new file mode 100755 index 000000000..dd66449b2 --- /dev/null +++ b/.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 + diff --git a/.github/workflows/aiter-prebuilt-upload.yml b/.github/workflows/aiter-prebuilt-upload.yml new file mode 100644 index 000000000..d9d9b856c --- /dev/null +++ b/.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 + ' 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