Update QoLA/AITER

transformer_engine/common/ck_fused_attn/CMakeLists.txt

@@ -65,48 +65,45 @@ if("${AITER_SHA}" STREQUAL "")
   message(FATAL_ERROR
           "Failed to parse 'aiter_commit = \"...\"' line in ${__QOLA_MANIFEST}.")
 endif()
-
-if(Python_EXECUTABLE)
-  if (NOT __SKIP_AITER_CHECKOUT)
-    execute_process(
-      COMMAND sh -c
-        "PYTHONPATH=\"${__QOLA_DIR}:$PYTHONPATH\" '${Python_EXECUTABLE}' -m qola.cli checkout \
-        --manifest '${__QOLA_MANIFEST}' \
-        --aiter-root '${__AITER_SOURCE_DIR}'"
-      RESULT_VARIABLE AITER_CHECKOUT_RESULT
-      OUTPUT_VARIABLE AITER_CHECKOUT_OUTPUT
-      ERROR_VARIABLE AITER_CHECKOUT_ERROR
-      OUTPUT_STRIP_TRAILING_WHITESPACE
-      ERROR_STRIP_TRAILING_WHITESPACE
-    )
-    if(NOT AITER_CHECKOUT_RESULT EQUAL 0)
-      message(FATAL_ERROR
-              "Failed to sync AITER source tree at ${__AITER_SOURCE_DIR} to "
-              "manifest-pinned commit ${AITER_SHA}.\n"
-              "${AITER_CHECKOUT_OUTPUT}\n${AITER_CHECKOUT_ERROR}")
-    endif()
-    message(STATUS "[AITER] Synced ${__AITER_SOURCE_DIR} to ${AITER_SHA}")
-  endif()
-
+if (NOT __SKIP_AITER_CHECKOUT)
   execute_process(
-    COMMAND ${Python_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/check_aiter_mha_args.py
-            --mode both
-            --te-dir "${CMAKE_CURRENT_LIST_DIR}/../../.."
-            --aiter-root "${__AITER_SOURCE_DIR}"
-    RESULT_VARIABLE AITER_ARG_CHECK_RESULT
-    OUTPUT_VARIABLE AITER_ARG_CHECK_OUTPUT
-    ERROR_VARIABLE AITER_ARG_CHECK_ERROR
+    COMMAND sh -c
+      "PYTHONPATH=\"${__QOLA_DIR}:$PYTHONPATH\" '${Python_EXECUTABLE}' -m qola.cli checkout \
+--manifest '${__QOLA_MANIFEST}' \
+--aiter-root '${__AITER_SOURCE_DIR}'"
+    RESULT_VARIABLE AITER_CHECKOUT_RESULT
+    OUTPUT_VARIABLE AITER_CHECKOUT_OUTPUT
+    ERROR_VARIABLE AITER_CHECKOUT_ERROR
     OUTPUT_STRIP_TRAILING_WHITESPACE
     ERROR_STRIP_TRAILING_WHITESPACE
   )
-
-  if(NOT AITER_ARG_CHECK_RESULT EQUAL 0)
+  if(NOT AITER_CHECKOUT_RESULT EQUAL 0)
     message(FATAL_ERROR
-            "AITER API validation failed in check_aiter_mha_args.py.\n"
-            "${AITER_ARG_CHECK_OUTPUT}\n${AITER_ARG_CHECK_ERROR}")
+            "Failed to sync AITER source tree at ${__AITER_SOURCE_DIR} to "
+            "manifest-pinned commit ${AITER_SHA}.\n"
+            "${AITER_CHECKOUT_OUTPUT}\n${AITER_CHECKOUT_ERROR}")
   endif()
-  message(STATUS "AITER API validation passed via check_aiter_mha_args.py")
+  message(STATUS "[AITER] Synced ${__AITER_SOURCE_DIR} to ${AITER_SHA}")
+endif()
+
+execute_process(
+  COMMAND ${Python_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/check_aiter_mha_args.py
+          --mode both
+          --te-dir "${CMAKE_CURRENT_LIST_DIR}/../../.."
+          --aiter-root "${__AITER_SOURCE_DIR}"
+  RESULT_VARIABLE AITER_ARG_CHECK_RESULT
+  OUTPUT_VARIABLE AITER_ARG_CHECK_OUTPUT
+  ERROR_VARIABLE AITER_ARG_CHECK_ERROR
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+  ERROR_STRIP_TRAILING_WHITESPACE
+)
+
+if(NOT AITER_ARG_CHECK_RESULT EQUAL 0)
+  message(FATAL_ERROR
+          "AITER API validation failed in check_aiter_mha_args.py.\n"
+          "${AITER_ARG_CHECK_OUTPUT}\n${AITER_ARG_CHECK_ERROR}")
 endif()
+message(STATUS "AITER API validation passed via check_aiter_mha_args.py")
 
 # Sanity-check the resolved include directories now that `qola checkout` has
 # materialized the AITER tree.
@@ -224,7 +221,8 @@ endforeach()
 add_library(ck_fused_attn SHARED ${ck_fused_attn_SOURCES})
 set(CK_FUSED_ATTN_COMPILE_OPTIONS)
 list(APPEND CK_FUSED_ATTN_COMPILE_OPTIONS
-    -DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=${CK_FUSED_ATTN_FLOAT_TO_BFLOAT16_DEFAULT})
+    -DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=${CK_FUSED_ATTN_FLOAT_TO_BFLOAT16_DEFAULT}
+    -DENABLE_CK=1)
 
 # Public QoLA headers ship alongside the .so libs in ${__AITER_MHA_PATH}/../include
 # (emitted by qola.cli build, or copied from the QoLA build dir above for the

transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp

@@ -113,7 +113,6 @@ struct CkAttnBwdArgs : CKAttnCommonArgs {
   // dQ
   void* dq_ptr = nullptr;
   uint64_t stride_b_dq = 0, stride_h_dq = 0, stride_s_dq = 0;
-  void* dq_acc_ptr = nullptr;
 
   // dK / dV expanded (MQA/GQA reduction inputs; null when h==hg)
   void* dk_expanded_ptr = nullptr;
@@ -134,6 +133,13 @@ struct CkAttnBwdArgs : CKAttnCommonArgs {
   // Workspace shared with forward LSE
   void* lse_workspace_ptr = nullptr;
 
+  // AOT scratch for AITER's internal bwd allocations (launcher metadata + dq_acc
+  // accumulator). Carved from the caller's workspace and handed to aiter through
+  // the workspace_alloc callback; ck_attn_bwd_workspace_size() reports the bytes
+  // to reserve. aiter_workspace_bytes bounds the bump allocator.
+  void* aiter_workspace_ptr = nullptr;
+  size_t aiter_workspace_bytes = 0;
+
   // V3 ASM kernel selection
   bool deterministic = false;
   bool uses_bwd_v3 = false;
@@ -143,6 +149,10 @@ struct CkAttnBwdArgs : CKAttnCommonArgs {
 hipError_t ck_attn_fwd(const CKAttnFwdArgs& args, hipStream_t stream);
 hipError_t ck_attn_bwd(const CkAttnBwdArgs& args, hipStream_t stream);
 
+// Bytes of AOT device scratch ck_attn_bwd needs for AITER's internal bwd
+// workspace (launcher metadata + dq_acc), covering both the v2 (CK launcher) and
+// v3 (asm) dispatch paths. Pure host-side computation; no kernel launch.
+size_t ck_attn_bwd_workspace_size(const CkAttnBwdArgs& args);
 // Probe whether AITER's v3 (asm) path will run for the given config, without
 // launching a kernel (backed by AITER's v3_api_check dry-run). Returns true iff
 // the v3 path is selected; false means the CK v2 path (or no support) would run.

transformer_engine/common/ck_fused_attn/qola_manifest.toml

@@ -1,5 +1,5 @@
 [qola]
-aiter_commit = "d32b0cb62ecc32bb1a858e8437d58eb9b3856af6"    # pinned AITER submodule commit
+aiter_commit = "e95f7d4e3924e848c1daf4df449e27822b332a09"    # pinned AITER submodule commit
 namespace = "te"
 rocm_versions = ["7.2"]
 
@@ -9,9 +9,11 @@ architectures = ["gfx950", "gfx942"]
 [[modules]]
 name = "libmha_fwd"
 mode = "cpp_itfs"
+receipt = 700
 drop_srcs = ["mha_fwd_split.cu", "mha_fwd_batch_prefill.cu"]
 drop_directions = ["fwd_splitkv", "batch_prefill"]
 
 [[modules]]
 name = "libmha_bwd"
 mode = "cpp_itfs"
+receipt = 700