From c9402a0cbcfdc1258e0d4165b4abf4a3abe5f774 Mon Sep 17 00:00:00 2001
From: Meekail Zain <zainmeekail@gmail.com>
Date: Mon, 2 Mar 2026 12:20:59 -0600
Subject: [PATCH 1/3] Added static build path for AITER backend

---
 setup.py                                      |  4 ++
 .../common/ck_fused_attn/CMakeLists.txt       | 22 +++++++--
 .../common/ck_fused_attn/aiter_build.sh       | 46 ++++++++++++++++++-
 .../common/ck_fused_attn/aiter_prebuilt.cmake |  3 ++
 4 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/setup.py b/setup.py
index 2f3e3c2ab..0c0894a3f 100644
--- a/setup.py
+++ b/setup.py
@@ -79,6 +79,10 @@ def setup_common_extension() -> CMakeExtension:
             ck_path = Path(os.getenv("NVTE_CK_FUSED_ATTN_PATH"))
             cmake_flags.append(f"-DAITER_MHA_PATH={ck_path}")
 
+        if os.getenv("NVTE_AITER_STATIC_LINK") is not None:
+            aiter_static_link = "ON" if int(os.getenv("NVTE_AITER_STATIC_LINK", "0")) else "OFF"
+            cmake_flags.append(f"-DNVTE_AITER_STATIC_LINK={aiter_static_link}")
+
         if int(os.getenv("NVTE_FUSED_ATTN_AOTRITON", "1"))==0 or int(os.getenv("NVTE_FUSED_ATTN", "1"))==0:
             cmake_flags.append("-DUSE_FUSED_ATTN_AOTRITON=OFF")
         elif os.getenv("NVTE_FUSED_ATTN_AOTRITON") or os.getenv("NVTE_FUSED_ATTN"):
diff --git a/transformer_engine/common/ck_fused_attn/CMakeLists.txt b/transformer_engine/common/ck_fused_attn/CMakeLists.txt
index 1c61596c0..742e67869 100644
--- a/transformer_engine/common/ck_fused_attn/CMakeLists.txt
+++ b/transformer_engine/common/ck_fused_attn/CMakeLists.txt
@@ -7,6 +7,7 @@ project(ck_fused_attn LANGUAGES HIP CXX)
 
 
 set(AITER_MHA_INSTALL_PREFIX "transformer_engine" CACHE STRING "aiter mha shared lib install prefix in TE")
+option(NVTE_AITER_STATIC_LINK "Statically link AITER MHA libs into ck_fused_attn" OFF)
 
 set(__AITER_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}/../../../3rdparty/aiter")
 set(__AITER_TEST_DIR "${__AITER_SOURCE_DIR}/op_tests/cpp/mha")
@@ -33,7 +34,7 @@ list(JOIN V3_ASM_ARCHS ";" V3_ASM_ARCHS_STR)
 
 if(DEFINED AITER_MHA_PATH)
   message(STATUS "[AITER-PREBUILT] Using AITER_MHA_PATH=${AITER_MHA_PATH}")
-  # use pre-built libmha_fwd.so libmha_bwd.so
+  # use pre-built libraries
   set(__AITER_MHA_PATH ${AITER_MHA_PATH})
 else()
   set(AITER_CACHE_VALID FALSE)
@@ -110,10 +111,25 @@ target_include_directories(ck_fused_attn PRIVATE ${CK_INCLUDE_DIR} ${__CK_SOURCE
 target_include_directories(ck_fused_attn PRIVATE ${AITER_INCLUDE_DIR})
 
 find_package(hip)
-list(APPEND ck_fused_attn_LINKER_LIBS hip::host hip::device roctx64 ${__AITER_MHA_PATH}/libmha_fwd.so ${__AITER_MHA_PATH}/libmha_bwd.so)
+if(NVTE_AITER_STATIC_LINK)
+  set(__AITER_MHA_FWD_LIB "${__AITER_MHA_PATH}/libmha_fwd.a")
+  set(__AITER_MHA_BWD_LIB "${__AITER_MHA_PATH}/libmha_bwd.a")
+  if(NOT EXISTS "${__AITER_MHA_FWD_LIB}" OR NOT EXISTS "${__AITER_MHA_BWD_LIB}")
+    message(FATAL_ERROR
+            "NVTE_AITER_STATIC_LINK=ON requires libmha_fwd.a and libmha_bwd.a under ${__AITER_MHA_PATH}. "
+            "Rebuild AITER with static archive generation enabled.")
+  endif()
+  list(APPEND ck_fused_attn_LINKER_LIBS hip::host hip::device roctx64 ${__AITER_MHA_FWD_LIB} ${__AITER_MHA_BWD_LIB})
+  target_link_options(ck_fused_attn PRIVATE -Wl,--exclude-libs,ALL)
+  message(STATUS "Statically linking AITER MHA libs into ck_fused_attn")
+else()
+  list(APPEND ck_fused_attn_LINKER_LIBS hip::host hip::device roctx64 ${__AITER_MHA_PATH}/libmha_fwd.so ${__AITER_MHA_PATH}/libmha_bwd.so)
+endif()
 target_link_libraries(ck_fused_attn PUBLIC ${ck_fused_attn_LINKER_LIBS})
 target_compile_options(ck_fused_attn PRIVATE ${CK_FUSED_ATTN_COMPILE_OPTIONS})
 set_target_properties(ck_fused_attn PROPERTIES INSTALL_RPATH "$ORIGIN")
 
-install(FILES ${__AITER_MHA_PATH}/libmha_fwd.so ${__AITER_MHA_PATH}/libmha_bwd.so DESTINATION ${CMAKE_INSTALL_PREFIX}/${AITER_MHA_INSTALL_PREFIX}/lib)
+if(NOT NVTE_AITER_STATIC_LINK)
+  install(FILES ${__AITER_MHA_PATH}/libmha_fwd.so ${__AITER_MHA_PATH}/libmha_bwd.so DESTINATION ${CMAKE_INSTALL_PREFIX}/${AITER_MHA_INSTALL_PREFIX}/lib)
+endif()
 install(TARGETS ck_fused_attn DESTINATION ${CMAKE_INSTALL_PREFIX}/${AITER_MHA_INSTALL_PREFIX}/lib)
diff --git a/transformer_engine/common/ck_fused_attn/aiter_build.sh b/transformer_engine/common/ck_fused_attn/aiter_build.sh
index 3ccf2979c..5c4114b50 100644
--- a/transformer_engine/common/ck_fused_attn/aiter_build.sh
+++ b/transformer_engine/common/ck_fused_attn/aiter_build.sh
@@ -32,7 +32,7 @@ while [[ $# -gt 0 ]]; do
 done
 
 if [[ -z "${AITER_DIR}" || -z "${AITER_TEST_DIR}" || -z "${GPU_ARCHS_VAL}" ]]; then
-  echo "[AITER-PREBUILT] --aiter-dir, --aiter-test-dir, and --gpu-archs are required." >&2
+  echo "[AITER-BUILD] --aiter-dir, --aiter-test-dir, and --gpu-archs are required." >&2
   exit 1
 fi
 
@@ -42,3 +42,47 @@ CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT="${CK_TILE_BF16_DEFAULT}" \
 GPU_ARCHS="${GPU_ARCHS_VAL}" \
 python3 "${AITER_TEST_DIR}/compile.py"
 
+# Generate static archives from the built object files only if NVTE_AITER_STATIC_LINK=1
+if [[ "${NVTE_AITER_STATIC_LINK:-0}" -ne 1 ]]; then
+  exit 0
+fi
+
+# Check for ar and ranlib
+AR_BIN="${AR:-$(command -v ar || true)}"
+RANLIB_BIN="${RANLIB:-$(command -v ranlib || true)}"
+if [[ -z "${AR_BIN}" ]]; then
+  echo "[AITER-BUILD] Could not find ar for static archive generation." >&2
+  exit 1
+fi
+if [[ -z "${RANLIB_BIN}" ]]; then
+  echo "[AITER-BUILD] Could not find ranlib for static archive generation." >&2
+  exit 1
+fi
+
+# Create static archives for both forward and backward passes
+for lib in fwd bwd; do
+  src_obj_dir="${AITER_DIR}/aiter/jit/build/libmha_${lib}/build"
+  out_archive="${AITER_TEST_DIR}/libmha_${lib}.a"
+
+  if [[ ! -d "${src_obj_dir}" ]]; then
+    echo "[AITER-BUILD] Missing object directory: ${src_obj_dir}" >&2
+    exit 1
+  fi
+
+  mapfile -d '' obj_files < <(find "${src_obj_dir}" -type f -name '*.o' -print0)
+  if [[ ${#obj_files[@]} -eq 0 ]]; then
+    echo "[AITER-BUILD] No object files found under ${src_obj_dir}" >&2
+    exit 1
+  fi
+
+  rm -f "${out_archive}"
+  for obj in "${obj_files[@]}"; do
+    "${AR_BIN}" q "${out_archive}" "${obj}"
+  done
+  if [[ -n "${RANLIB_BIN}" ]]; then
+    "${RANLIB_BIN}" "${out_archive}"
+  fi
+
+  echo "[AITER-BUILD] Created static archive: ${out_archive} (${#obj_files[@]} objects)"
+done
+
diff --git a/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake b/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake
index a59605e00..a9ff3fe4d 100644
--- a/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake
+++ b/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake
@@ -45,6 +45,9 @@ function(cache_local_aiter_build SOURCE_DIR)
   file(MAKE_DIRECTORY "${EXTRACT_DIR}")
   message(STATUS "[AITER-PREBUILT] Caching locally built libs to ${EXTRACT_DIR}")
   file(COPY "${SOURCE_DIR}/libmha_fwd.so" "${SOURCE_DIR}/libmha_bwd.so" DESTINATION "${EXTRACT_DIR}")
+  if(NVTE_AITER_STATIC_LINK)
+    file(COPY "${SOURCE_DIR}/libmha_fwd.a" "${SOURCE_DIR}/libmha_bwd.a" DESTINATION "${EXTRACT_DIR}")
+  endif()
 endfunction()
 
 # Download prebuilt tgz file

From 2c4c11bd35ca97e833009b32c7f77ebafb4083e1 Mon Sep 17 00:00:00 2001
From: Meekail Zain <zainmeekail@gmail.com>
Date: Mon, 2 Mar 2026 15:40:26 -0600
Subject: [PATCH 2/3] Updated w/ batched progress tracking for archiving

---
 setup.py                                      |  2 +-
 .../common/ck_fused_attn/aiter_build.sh       | 19 ++++++++++++++--
 .../common/ck_fused_attn/aiter_prebuilt.cmake | 22 +++++++++++++++++--
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/setup.py b/setup.py
index 0c0894a3f..2770e093f 100644
--- a/setup.py
+++ b/setup.py
@@ -80,7 +80,7 @@ def setup_common_extension() -> CMakeExtension:
             cmake_flags.append(f"-DAITER_MHA_PATH={ck_path}")
 
         if os.getenv("NVTE_AITER_STATIC_LINK") is not None:
-            aiter_static_link = "ON" if int(os.getenv("NVTE_AITER_STATIC_LINK", "0")) else "OFF"
+            aiter_static_link = "ON" if int(os.getenv("NVTE_AITER_STATIC_LINK", "1")) else "OFF"
             cmake_flags.append(f"-DNVTE_AITER_STATIC_LINK={aiter_static_link}")
 
         if int(os.getenv("NVTE_FUSED_ATTN_AOTRITON", "1"))==0 or int(os.getenv("NVTE_FUSED_ATTN", "1"))==0:
diff --git a/transformer_engine/common/ck_fused_attn/aiter_build.sh b/transformer_engine/common/ck_fused_attn/aiter_build.sh
index 5c4114b50..40ef9a259 100644
--- a/transformer_engine/common/ck_fused_attn/aiter_build.sh
+++ b/transformer_engine/common/ck_fused_attn/aiter_build.sh
@@ -60,6 +60,7 @@ if [[ -z "${RANLIB_BIN}" ]]; then
 fi
 
 # Create static archives for both forward and backward passes
+batch_size=${AITER_ARCHIVE_BATCH_SIZE:-400}
 for lib in fwd bwd; do
   src_obj_dir="${AITER_DIR}/aiter/jit/build/libmha_${lib}/build"
   out_archive="${AITER_TEST_DIR}/libmha_${lib}.a"
@@ -75,9 +76,23 @@ for lib in fwd bwd; do
     exit 1
   fi
 
+  total_objs=${#obj_files[@]}
+
   rm -f "${out_archive}"
-  for obj in "${obj_files[@]}"; do
-    "${AR_BIN}" q "${out_archive}" "${obj}"
+  idx=0
+  while [[ ${idx} -lt ${total_objs} ]]; do
+    remaining=$((total_objs - idx))
+    current_batch_size=${batch_size}
+    if [[ ${remaining} -lt ${batch_size} ]]; then
+      current_batch_size=${remaining}
+    fi
+
+    chunk=("${obj_files[@]:idx:current_batch_size}")
+    "${AR_BIN}" q "${out_archive}" "${chunk[@]}"
+
+    idx=$((idx + current_batch_size))
+    percent=$((idx * 100 / total_objs))
+    echo -ne "[AITER-BUILD] [${lib}] archiving objects: ${idx}/${total_objs} (${percent}%)\r"
   done
   if [[ -n "${RANLIB_BIN}" ]]; then
     "${RANLIB_BIN}" "${out_archive}"
diff --git a/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake b/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake
index a9ff3fe4d..4ee397b45 100644
--- a/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake
+++ b/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake
@@ -29,7 +29,18 @@ set(EXTRACT_DIR "${CACHE_ROOT}/${KEY}")
 
 # Validate existing cache path
 function(is_aiter_cache_valid CACHE_VALID)
-  if(EXISTS "${EXTRACT_DIR}/libmha_fwd.so" AND EXISTS "${EXTRACT_DIR}/libmha_bwd.so")
+  set(_AITER_CACHE_VALID TRUE)
+  if(NVTE_AITER_STATIC_LINK)
+    set(_AITER_LIB_EXT "a")
+  else()
+    set(_AITER_LIB_EXT "so")
+  endif()
+
+  if(NOT (EXISTS "${EXTRACT_DIR}/libmha_fwd.${_AITER_LIB_EXT}" AND EXISTS "${EXTRACT_DIR}/libmha_bwd.${_AITER_LIB_EXT}"))
+    set(_AITER_CACHE_VALID FALSE)
+  endif()
+
+  if(_AITER_CACHE_VALID)
     set(${CACHE_VALID} TRUE PARENT_SCOPE)
     message(STATUS "[AITER-PREBUILT] Found Cached build files at ${EXTRACT_DIR}")
     return()
@@ -44,8 +55,15 @@ endfunction()
 function(cache_local_aiter_build SOURCE_DIR)
   file(MAKE_DIRECTORY "${EXTRACT_DIR}")
   message(STATUS "[AITER-PREBUILT] Caching locally built libs to ${EXTRACT_DIR}")
-  file(COPY "${SOURCE_DIR}/libmha_fwd.so" "${SOURCE_DIR}/libmha_bwd.so" DESTINATION "${EXTRACT_DIR}")
+  if(EXISTS "${SOURCE_DIR}/libmha_fwd.so" AND EXISTS "${SOURCE_DIR}/libmha_bwd.so")
+    file(COPY "${SOURCE_DIR}/libmha_fwd.so" "${SOURCE_DIR}/libmha_bwd.so" DESTINATION "${EXTRACT_DIR}")
+  elseif(NOT NVTE_AITER_STATIC_LINK)
+    message(FATAL_ERROR "Expected libmha_fwd.so and libmha_bwd.so under ${SOURCE_DIR}")
+  endif()
   if(NVTE_AITER_STATIC_LINK)
+    if(NOT EXISTS "${SOURCE_DIR}/libmha_fwd.a" OR NOT EXISTS "${SOURCE_DIR}/libmha_bwd.a")
+      message(FATAL_ERROR "Expected libmha_fwd.a and libmha_bwd.a under ${SOURCE_DIR} for static link mode")
+    endif()
     file(COPY "${SOURCE_DIR}/libmha_fwd.a" "${SOURCE_DIR}/libmha_bwd.a" DESTINATION "${EXTRACT_DIR}")
   endif()
 endfunction()

From 95c2f0544407327ca6c6303079ab64dbdea3445e Mon Sep 17 00:00:00 2001
From: Meekail Zain <zainmeekail@gmail.com>
Date: Tue, 3 Mar 2026 11:26:14 -0600
Subject: [PATCH 3/3] Updated cmake var name, streamlined implementation

---
 setup.py                                      |  2 +-
 .../common/ck_fused_attn/CMakeLists.txt       | 31 ++++++++++---------
 .../common/ck_fused_attn/aiter_build.sh       | 28 +++++------------
 .../common/ck_fused_attn/aiter_prebuilt.cmake | 12 +++----
 4 files changed, 28 insertions(+), 45 deletions(-)

diff --git a/setup.py b/setup.py
index 2770e093f..91e522f40 100644
--- a/setup.py
+++ b/setup.py
@@ -81,7 +81,7 @@ def setup_common_extension() -> CMakeExtension:
 
         if os.getenv("NVTE_AITER_STATIC_LINK") is not None:
             aiter_static_link = "ON" if int(os.getenv("NVTE_AITER_STATIC_LINK", "1")) else "OFF"
-            cmake_flags.append(f"-DNVTE_AITER_STATIC_LINK={aiter_static_link}")
+            cmake_flags.append(f"-DAITER_STATIC_LINK={aiter_static_link}")
 
         if int(os.getenv("NVTE_FUSED_ATTN_AOTRITON", "1"))==0 or int(os.getenv("NVTE_FUSED_ATTN", "1"))==0:
             cmake_flags.append("-DUSE_FUSED_ATTN_AOTRITON=OFF")
diff --git a/transformer_engine/common/ck_fused_attn/CMakeLists.txt b/transformer_engine/common/ck_fused_attn/CMakeLists.txt
index 742e67869..cbe7ae7ef 100644
--- a/transformer_engine/common/ck_fused_attn/CMakeLists.txt
+++ b/transformer_engine/common/ck_fused_attn/CMakeLists.txt
@@ -7,7 +7,7 @@ project(ck_fused_attn LANGUAGES HIP CXX)
 
 
 set(AITER_MHA_INSTALL_PREFIX "transformer_engine" CACHE STRING "aiter mha shared lib install prefix in TE")
-option(NVTE_AITER_STATIC_LINK "Statically link AITER MHA libs into ck_fused_attn" OFF)
+option(AITER_STATIC_LINK "Statically link AITER MHA libs into ck_fused_attn" OFF)
 
 set(__AITER_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}/../../../3rdparty/aiter")
 set(__AITER_TEST_DIR "${__AITER_SOURCE_DIR}/op_tests/cpp/mha")
@@ -33,7 +33,7 @@ message(STATUS "AITER V3_ASM_ARCHS: ${V3_ASM_ARCHS}")
 list(JOIN V3_ASM_ARCHS ";" V3_ASM_ARCHS_STR)
 
 if(DEFINED AITER_MHA_PATH)
-  message(STATUS "[AITER-PREBUILT] Using AITER_MHA_PATH=${AITER_MHA_PATH}")
+  message(STATUS "[AITER-BUILD] Using AITER_MHA_PATH=${AITER_MHA_PATH}")
   # use pre-built libraries
   set(__AITER_MHA_PATH ${AITER_MHA_PATH})
 else()
@@ -50,7 +50,7 @@ else()
     
     # If not downloaded, Fallback: Build from source
     if(NOT AITER_PREBUILT_DOWNLOAD_SUCCESS)
-      message(STATUS " [AITER-PREBUILT] Building aiter from source.")
+      message(STATUS " [AITER-BUILD] Building aiter from source.")
       execute_process(
         COMMAND bash ${CMAKE_CURRENT_LIST_DIR}/aiter_build.sh
                 --aiter-dir ${__AITER_SOURCE_DIR}
@@ -63,7 +63,7 @@ else()
     endif()
   endif()
   set(__AITER_MHA_PATH "${EXTRACT_DIR}")
-  message(STATUS "[AITER-PREBUILT] Using __AITER_MHA_PATH=${__AITER_MHA_PATH}")
+  message(STATUS "[AITER-BUILD] Using __AITER_MHA_PATH=${__AITER_MHA_PATH}")
 endif()
 
 set(ck_fused_attn_SOURCES)
@@ -111,25 +111,26 @@ target_include_directories(ck_fused_attn PRIVATE ${CK_INCLUDE_DIR} ${__CK_SOURCE
 target_include_directories(ck_fused_attn PRIVATE ${AITER_INCLUDE_DIR})
 
 find_package(hip)
-if(NVTE_AITER_STATIC_LINK)
-  set(__AITER_MHA_FWD_LIB "${__AITER_MHA_PATH}/libmha_fwd.a")
-  set(__AITER_MHA_BWD_LIB "${__AITER_MHA_PATH}/libmha_bwd.a")
-  if(NOT EXISTS "${__AITER_MHA_FWD_LIB}" OR NOT EXISTS "${__AITER_MHA_BWD_LIB}")
-    message(FATAL_ERROR
-            "NVTE_AITER_STATIC_LINK=ON requires libmha_fwd.a and libmha_bwd.a under ${__AITER_MHA_PATH}. "
-            "Rebuild AITER with static archive generation enabled.")
-  endif()
-  list(APPEND ck_fused_attn_LINKER_LIBS hip::host hip::device roctx64 ${__AITER_MHA_FWD_LIB} ${__AITER_MHA_BWD_LIB})
+list(APPEND ck_fused_attn_LINKER_LIBS hip::host hip::device roctx64)
+
+if(AITER_STATIC_LINK)
+  set(_AITER_LIB_EXT "a")
   target_link_options(ck_fused_attn PRIVATE -Wl,--exclude-libs,ALL)
   message(STATUS "Statically linking AITER MHA libs into ck_fused_attn")
 else()
-  list(APPEND ck_fused_attn_LINKER_LIBS hip::host hip::device roctx64 ${__AITER_MHA_PATH}/libmha_fwd.so ${__AITER_MHA_PATH}/libmha_bwd.so)
+  set(_AITER_LIB_EXT "so")
+  message(STATUS "Using dynamic AITER MHA libs for ck_fused_attn")
 endif()
+
+set(__AITER_MHA_FWD_LIB "${__AITER_MHA_PATH}/libmha_fwd.${_AITER_LIB_EXT}")
+set(__AITER_MHA_BWD_LIB "${__AITER_MHA_PATH}/libmha_bwd.${_AITER_LIB_EXT}")
+list(APPEND ck_fused_attn_LINKER_LIBS ${__AITER_MHA_FWD_LIB} ${__AITER_MHA_BWD_LIB})
+
 target_link_libraries(ck_fused_attn PUBLIC ${ck_fused_attn_LINKER_LIBS})
 target_compile_options(ck_fused_attn PRIVATE ${CK_FUSED_ATTN_COMPILE_OPTIONS})
 set_target_properties(ck_fused_attn PROPERTIES INSTALL_RPATH "$ORIGIN")
 
-if(NOT NVTE_AITER_STATIC_LINK)
+if(NOT AITER_STATIC_LINK)
   install(FILES ${__AITER_MHA_PATH}/libmha_fwd.so ${__AITER_MHA_PATH}/libmha_bwd.so DESTINATION ${CMAKE_INSTALL_PREFIX}/${AITER_MHA_INSTALL_PREFIX}/lib)
 endif()
 install(TARGETS ck_fused_attn DESTINATION ${CMAKE_INSTALL_PREFIX}/${AITER_MHA_INSTALL_PREFIX}/lib)
diff --git a/transformer_engine/common/ck_fused_attn/aiter_build.sh b/transformer_engine/common/ck_fused_attn/aiter_build.sh
index 40ef9a259..c8b96afbb 100644
--- a/transformer_engine/common/ck_fused_attn/aiter_build.sh
+++ b/transformer_engine/common/ck_fused_attn/aiter_build.sh
@@ -36,14 +36,14 @@ if [[ -z "${AITER_DIR}" || -z "${AITER_TEST_DIR}" || -z "${GPU_ARCHS_VAL}" ]]; t
   exit 1
 fi
 
-rm -rf "${AITER_DIR}/aiter/jit/build"
-AITER_LOG_MORE=1 \
-CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT="${CK_TILE_BF16_DEFAULT}" \
-GPU_ARCHS="${GPU_ARCHS_VAL}" \
-python3 "${AITER_TEST_DIR}/compile.py"
+# rm -rf "${AITER_DIR}/aiter/jit/build"
+# AITER_LOG_MORE=1 \
+# CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT="${CK_TILE_BF16_DEFAULT}" \
+# GPU_ARCHS="${GPU_ARCHS_VAL}" \
+# python3 "${AITER_TEST_DIR}/compile.py"
 
 # Generate static archives from the built object files only if NVTE_AITER_STATIC_LINK=1
-if [[ "${NVTE_AITER_STATIC_LINK:-0}" -ne 1 ]]; then
+if [[ "${NVTE_AITER_STATIC_LINK:-1}" -ne 1 ]]; then
   exit 0
 fi
 
@@ -60,7 +60,6 @@ if [[ -z "${RANLIB_BIN}" ]]; then
 fi
 
 # Create static archives for both forward and backward passes
-batch_size=${AITER_ARCHIVE_BATCH_SIZE:-400}
 for lib in fwd bwd; do
   src_obj_dir="${AITER_DIR}/aiter/jit/build/libmha_${lib}/build"
   out_archive="${AITER_TEST_DIR}/libmha_${lib}.a"
@@ -79,21 +78,8 @@ for lib in fwd bwd; do
   total_objs=${#obj_files[@]}
 
   rm -f "${out_archive}"
-  idx=0
-  while [[ ${idx} -lt ${total_objs} ]]; do
-    remaining=$((total_objs - idx))
-    current_batch_size=${batch_size}
-    if [[ ${remaining} -lt ${batch_size} ]]; then
-      current_batch_size=${remaining}
-    fi
+  "${AR_BIN}" q "${out_archive}" "${obj_files[@]}"
 
-    chunk=("${obj_files[@]:idx:current_batch_size}")
-    "${AR_BIN}" q "${out_archive}" "${chunk[@]}"
-
-    idx=$((idx + current_batch_size))
-    percent=$((idx * 100 / total_objs))
-    echo -ne "[AITER-BUILD] [${lib}] archiving objects: ${idx}/${total_objs} (${percent}%)\r"
-  done
   if [[ -n "${RANLIB_BIN}" ]]; then
     "${RANLIB_BIN}" "${out_archive}"
   fi
diff --git a/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake b/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake
index 4ee397b45..b19a3fa72 100644
--- a/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake
+++ b/transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake
@@ -30,7 +30,7 @@ set(EXTRACT_DIR "${CACHE_ROOT}/${KEY}")
 # Validate existing cache path
 function(is_aiter_cache_valid CACHE_VALID)
   set(_AITER_CACHE_VALID TRUE)
-  if(NVTE_AITER_STATIC_LINK)
+  if(AITER_STATIC_LINK)
     set(_AITER_LIB_EXT "a")
   else()
     set(_AITER_LIB_EXT "so")
@@ -47,7 +47,7 @@ function(is_aiter_cache_valid CACHE_VALID)
   endif()
 
   # Cache is invalid/outdated - clean it
-  file(REMOVE_RECURSE "${CACHE_ROOT}")
+  file(REMOVE_RECURSE "${EXTRACT_DIR}")
   file(REMOVE_RECURSE "${CMAKE_BINARY_DIR}/_deps")
 endfunction()
 
@@ -55,12 +55,8 @@ endfunction()
 function(cache_local_aiter_build SOURCE_DIR)
   file(MAKE_DIRECTORY "${EXTRACT_DIR}")
   message(STATUS "[AITER-PREBUILT] Caching locally built libs to ${EXTRACT_DIR}")
-  if(EXISTS "${SOURCE_DIR}/libmha_fwd.so" AND EXISTS "${SOURCE_DIR}/libmha_bwd.so")
-    file(COPY "${SOURCE_DIR}/libmha_fwd.so" "${SOURCE_DIR}/libmha_bwd.so" DESTINATION "${EXTRACT_DIR}")
-  elseif(NOT NVTE_AITER_STATIC_LINK)
-    message(FATAL_ERROR "Expected libmha_fwd.so and libmha_bwd.so under ${SOURCE_DIR}")
-  endif()
-  if(NVTE_AITER_STATIC_LINK)
+  file(COPY "${SOURCE_DIR}/libmha_fwd.so" "${SOURCE_DIR}/libmha_bwd.so" DESTINATION "${EXTRACT_DIR}")
+  if(AITER_STATIC_LINK)
     if(NOT EXISTS "${SOURCE_DIR}/libmha_fwd.a" OR NOT EXISTS "${SOURCE_DIR}/libmha_bwd.a")
       message(FATAL_ERROR "Expected libmha_fwd.a and libmha_bwd.a under ${SOURCE_DIR} for static link mode")
     endif()