Updated w/ batched progress tracking for archiving

Author: Micky774

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}"

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()