Ipanfilo/aiter split kv fwd

tests/pytorch/attention/test_attention.py

@@ -134,6 +134,13 @@ def reset_attn_backend():
     "base_5_1": ModelConfig(8, 128, 16, 512, max_seqlen_kv=2048),
     "base_6_0": ModelConfig(8, 1, 16, 1024, max_seqlen_kv=2048),
     "base_6_1": ModelConfig(8, 128, 16, 1024, max_seqlen_kv=2048),
+    # gfx942 split-KV testing
+    "base_7_0": ModelConfig(1, 1, 8, 64, max_seqlen_kv=2048, num_gqa_groups=1),
+    "base_7_1": ModelConfig(1, 1, 8, 64, max_seqlen_kv=2048, num_gqa_groups=1,
+                            attn_mask_type="causal_bottom_right"),
+    #test fallthrough to bactch kernels with causal mask
+    "base_7_2": ModelConfig(1, 1, 8, 64, max_seqlen_kv=2048, num_gqa_groups=1,
+                            attn_mask_type="causal"),
 }
 
 
@@ -908,6 +915,9 @@ def test_dpa_alibi_slopes(dtype, model_configs, model):
         attn_mask_type="padding_causal",
         attn_bias_type="post_scale_bias",
     ),
+    "layout_0_4": ModelConfig(
+        2, 128, 16, 64, max_seqlen_kv=512, num_gqa_groups=1, attn_mask_type="causal"
+    ),
     "layout_1_0": ModelConfig(2, 2048, 24, 128),
     "layout_1_1": ModelConfig(
         2, 2048, 24, 128, attn_mask_type="causal", attn_bias_type="post_scale_bias"

transformer_engine/common/ck_fused_attn/CMakeLists.txt

@@ -66,13 +66,27 @@ if("${AITER_SHA}" STREQUAL "")
           "Failed to parse 'aiter_commit = \"...\"' line in ${__QOLA_MANIFEST}.")
 endif()
 
+# Resolve the manifest-pinned AITER commit (defines AITER_SHA) and bring the
+# QoLA-managed AITER source tree to that commit before any consumer reads it
+# (header validation below, header includes for the .cpp build later, and
+# QoLA's own kernel build if the prebuilt cache misses).
+include("${CMAKE_CURRENT_LIST_DIR}/aiter_prebuilt.cmake")
+
 if(Python_EXECUTABLE)
   if (NOT __SKIP_AITER_CHECKOUT)
+    # Redirect GIT_CONFIG_GLOBAL to a tempfile carrying `safe.directory = *` so
+    # git operations inside the QoLA-managed AITER tree (and its recursive
+    # submodules) work in containerized builds where the bind-mounted .git is
+    # owned by a different UID than the build process. Mirrors the pattern in
+    # transformer_engine/common/CMakeLists.txt:get_git_commit().
     execute_process(
       COMMAND sh -c
-        "PYTHONPATH=\"${__QOLA_DIR}:$PYTHONPATH\" '${Python_EXECUTABLE}' -m qola.cli checkout \
-        --manifest '${__QOLA_MANIFEST}' \
-        --aiter-root '${__AITER_SOURCE_DIR}'"
+        "tmp=$(mktemp /tmp/gitconfig.XXXXXX) || exit 1; \
+GIT_CONFIG_GLOBAL=$tmp git config --global --add safe.directory '*' >/dev/null 2>&1; \
+GIT_CONFIG_GLOBAL=$tmp PYTHONPATH=\"${__QOLA_DIR}:$PYTHONPATH\" '${Python_EXECUTABLE}' -m qola.cli checkout \
+--manifest '${__QOLA_MANIFEST}' \
+--aiter-root '${__AITER_SOURCE_DIR}'; \
+rc=$?; rm -f \"$tmp\"; exit $rc"
       RESULT_VARIABLE AITER_CHECKOUT_RESULT
       OUTPUT_VARIABLE AITER_CHECKOUT_OUTPUT
       ERROR_VARIABLE AITER_CHECKOUT_ERROR
@@ -106,6 +120,8 @@ if(Python_EXECUTABLE)
             "${AITER_ARG_CHECK_OUTPUT}\n${AITER_ARG_CHECK_ERROR}")
   endif()
   message(STATUS "AITER API validation passed via check_aiter_mha_args.py")
+else()
+  message(WARNING "Python interpreter not found; skipping AITER source-tree sync and API validation.")
 endif()
 
 # Sanity-check the resolved include directories now that `qola checkout` has
@@ -224,7 +240,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,18 @@ 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);
+
+// Reserve, outside any HIP graph capture, the pinned host staging buffer the v2 backward
+// launcher requests via ck_attn_bwd's pinned_host_alloc callback. Must be called only from
+// the workspace-size (sizing) pass, which runs at lowering time before stream capture;
+// allocating there keeps the captured dispatch allocation-free. Idempotent per config and
+// harmless for configs that end up on the v3 asm path (the buffer simply goes unused).
+void ck_attn_bwd_reserve_host_staging(const CkAttnBwdArgs& args);
+
 }//namespace ck_fused_attn
 #endif // CK_FUSED_ATTN_H
 

transformer_engine/common/ck_fused_attn/qola_manifest.toml

@@ -1,5 +1,5 @@
 [qola]
-aiter_commit = "d32b0cb62ecc32bb1a858e8437d58eb9b3856af6"    # pinned AITER submodule commit
+aiter_commit = "5ebc4cad3720c49fd2ee4c3ce41e9c2093070cd5"    # 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