Update AITER subcommit and refactor internal AITER/CK FA API usage

.gitignore

@@ -55,3 +55,4 @@ artifacts/
 **/times.csv
 transformer_engine/build_info.txt
 transformer_engine/common/util/hip_nvml.*
+transformer_engine/lib/aiter

setup.py

@@ -8,6 +8,7 @@
 
 from importlib import metadata
 import os
+import sys
 import time
 from pathlib import Path
 from typing import List, Tuple
@@ -88,6 +89,17 @@ def setup_common_extension() -> CMakeExtension:
             cmake_flags.append("-DUSE_FUSED_ATTN_CK=OFF")
         elif os.getenv("NVTE_FUSED_ATTN_CK") or os.getenv("NVTE_FUSED_ATTN"):
             cmake_flags.append("-DUSE_FUSED_ATTN_CK=ON")
+            # Dynamically set AITER_LOG_MORE based on PIP_VERBOSE to avoid excessive logging during build.
+            os.environ["AITER_LOG_MORE"] = str(max(int(os.environ.get('PIP_VERBOSE', '0')) - 1, 0))
+            # Explicitly checks the AITER API usage
+            try:
+                subprocess.run(
+                    sys.executable + " tools/check_aiter_mha_args_usage.py --mode both",
+                    shell=True, check=True
+                )
+            except subprocess.CalledProcessError:
+                print("Error checking AITER mha_args usage.")
+                sys.exit(1)
 
         if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", "0"))) and os.getenv("NVTE_ENABLE_ROCSHMEM") is None:
             os.environ["NVTE_ENABLE_ROCSHMEM"] = '1'

tools/check_aiter_mha_args_usage.py

@@ -0,0 +1,97 @@
+import argparse
+import re
+from pathlib import Path
+from typing import List, Set
+import sys
+
+def parse_with_skip_comments(buffer, line, regex, outputs):
+    # skip comments
+    stripped = line.strip()
+    if not stripped or stripped.startswith("//"):
+        return
+    line_no_comment = re.sub(r"//.*", "", line)
+    buffer[0] += " " + line_no_comment.strip()
+    if ";" not in line_no_comment:
+        return
+    match = regex.search(buffer[0])
+    if match:
+        outputs.append(match.group(1))
+    buffer[0] = ""
+
+
+def extract_fields_from_header(text: str, struct_name: str) -> List[str]:
+    struct_field_re = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)\s*(?:=[^;]*)?;\s*$")
+    struct_end_re = re.compile(r"^\s*};\s*$")
+
+    struct_start_re = re.compile(rf"\bstruct\s+{re.escape(struct_name)}\b")
+    lines = text.splitlines()
+    in_struct = False
+    fields: List[str] = []
+    buffer = [""]
+    for line in lines:
+        if not in_struct:
+            if struct_start_re.search(line):
+                in_struct = True
+            continue
+        if struct_end_re.search(line):
+            break
+        parse_with_skip_comments(buffer, line, struct_field_re, fields)
+    return fields
+
+
+def extract_usage_from_source(text: str, var_name: str) -> Set[str]:
+    assign_re = re.compile(rf"\b{re.escape(var_name)}\.([A-Za-z_][A-Za-z0-9_]*)\b\s*=")
+    assignments = []
+    lines = text.splitlines()
+    buffer = [""]
+    for line in lines:
+        parse_with_skip_comments(buffer, line, assign_re, assignments)
+    return set(assignments)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Check aiter args usage vs header definition")
+    parser.add_argument("--mode", choices=["fwd", "bwd", "both"], required=True, help="Mode: fwd, bwd, or both")
+    args = parser.parse_args()
+    modes = ["fwd", "bwd"] if args.mode == "both" else [args.mode]
+    mismatch = 0
+    for mode in modes:
+        header_path = Path(f"3rdparty/aiter/csrc/include/mha_{mode}.h")
+        source_path = Path(f"transformer_engine/common/ck_fused_attn/src/ck_fused_attn_{mode}.cpp")
+        header_text = header_path.read_text(encoding="utf-8")
+        source_text = source_path.read_text(encoding="utf-8")
+
+        header_fields = extract_fields_from_header(header_text, f"mha_{mode}_args")
+        header_set = set(header_fields)
+        used_fields = extract_usage_from_source(source_text, f"fmha_args")
+
+        missing_in_usage = sorted(header_set - used_fields)
+        unknown_in_header = sorted(used_fields - header_set)
+        mismatch += len(missing_in_usage) + len(unknown_in_header)
+
+        print(f"\nAnalyzing mha_{mode}_args\n")
+        print(f"mha_{mode}_args fields in header:", len(header_set))
+        print(f"mha_{mode}_args fields referenced in source:", len(used_fields))
+
+        if missing_in_usage:
+            print("\nFields present in header but not referenced in source:")
+            for name in missing_in_usage:
+                print(f"  - {name}")
+        else:
+            print("\nAll header fields are referenced in source.")
+
+        if unknown_in_header:
+            print("\nFields referenced in source but not in header:")
+            for name in unknown_in_header:
+                print(f"  - {name}")
+        else:
+            print("\nNo unknown fields referenced in source.")
+
+    if mismatch:
+        print(f"\nTotal mismatched fields: {mismatch}")
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

transformer_engine/common/ck_fused_attn/aiter_build.sh

@@ -37,7 +37,6 @@ if [[ -z "${AITER_DIR}" || -z "${AITER_TEST_DIR}" || -z "${GPU_ARCHS_VAL}" ]]; t
 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"