Updated CK/AITER Cmake Build

[Micky774]

Description

Updated the cmake build for CK/AITER FA backend to include QoLA/ck-jit builds as a custom command for proper inclusion in build graph, including proper dependency graphing and caching.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[github-actions]

Claude Walkthrough

Intent. The CK/AITER fused-attention kernel source-build (te_libmha_{fwd,bwd}.so) was previously invoked eagerly via execute_process at CMake configure time, so it re-ran on every reconfigure — i.e. every pip install, even when nothing relevant had changed. This PR moves the build into a proper add_custom_command/add_custom_target, letting Ninja rebuild the kernels only when a declared input actually changes.

Key changes.

• Introduces a ck_fused_attn_kernels build-time custom target that produces te_libmha_fwd.so / te_libmha_bwd.so under ${CMAKE_CURRENT_BINARY_DIR}/ck_kernels/lib (transformer_engine/common/ck_fused_attn/CMakeLists.txt:236).
• Converts both source-build paths — CK-JIT (CMakeLists.txt:184) and plain QoLA (CMakeLists.txt:212) — from execute_process + FATAL_ERROR-on-nonzero into add_custom_command invocations with proper DEPENDS, COMMENT, VERBATIM, and USES_TERMINAL.
• Adds a generated cache-key file ck_build_key.txt via file(GENERATE) that encodes NVTE_CK_JIT, GPU archs, the bf16 default, CK_JIT_EXTRA_CACHE_KEY, and AITER_SHA (CMakeLists.txt:172). Content-only writes make it a fine-grained input to the custom command's DEPENDS.
• For the CK-JIT path, globs orchestrator scripts/patches (*.py|*.sh|*.cpp|*.patch under ${__CK_JIT_SOURCE_DIR}) and adds them to DEPENDS so edits to the JIT driver retrigger the build (CMakeLists.txt:184).
• Wires the ck_fused_attn shared library to depend on ck_fused_attn_kernels via add_dependencies and installs the freshly built .sos (plus an OPTIONAL ck_jit/ runtime tree for the CK-JIT path) from the build tree instead of the prior in-place location (CMakeLists.txt:289).
• Relaxes the header existence check: qola_config.h is no longer required at configure time when we are source-building, since it is produced by the custom command (CMakeLists.txt:271).

Walkthrough.
The gating variable __CK_KERNELS_TARGET is new and is what everything downstream keys off: empty means "consuming prebuilt or user-supplied libs, nothing to build," non-empty means "we emitted a build-time target and ck_fused_attn must depend on it." That lets the prebuilt/AITER_MHA_PATH branch keep its existing install semantics untouched while the source-build branches funnel through the new custom-target machinery.

Both source-build branches now write their outputs to a single canonical location (${CMAKE_CURRENT_BINARY_DIR}/ck_kernels/lib) and export headers under a sibling include/. The QoLA branch additionally stages libs+headers into the aiter prebuilt cache (via get_default_aiter_cache_dir) as COMMAND steps in the same custom command — replacing the previous file(COPY …) calls that ran at configure time and would silently re-copy on every reconfigure. The CK-JIT branch skips the cache population because CK-JIT is expected to always build from source anyway.

The install() rules are split into a source-build arm (install from ${__CK_FWD_SO}/${__CK_BWD_SO} in the build tree, plus optional ck_jit/ runtime) and the existing prebuilt arm (install from __AITER_MHA_PATH), selected by whether __CK_KERNELS_TARGET is set. install(TARGETS ck_fused_attn …) is unchanged.

Testing. No tests added — this is a pure build-system change. Verification is implicit: an incremental pip install should now skip the CK kernel rebuild when only unrelated sources changed, and a change to AITER_SHA, GPU archs, tunables, or any CK-JIT script/patch should retrigger it via the ck_build_key.txt or script glob.

Notes for reviewers.

• The file(GLOB) on JIT scripts is evaluated at configure time, so newly added script files won't be picked up until the next cmake reconfigure. Given how rarely those files change, that seems acceptable, but worth noting.
• USES_TERMINAL on both custom commands is deliberate — kernel builds are long-running and want a live TTY under Ninja rather than buffered output.
• The prebuilt/AITER_MHA_PATH code path is behaviorally unchanged; only the source-build path is restructured.
• The QoLA-branch cache population now happens as copy_directory steps inside the custom command; previously it happened via file(COPY) at configure time. If anything downstream was inspecting the aiter cache during configure of a different subproject, that ordering has shifted to build time.

---

Generated by Claude. To request a code review, comment /claude review.

[github-actions]

Reviewed the CK/AITER cmake refactor that moves the QoLA / CK-JIT kernel build from configure-time execute_process to a build-time add_custom_command fronted by an add_custom_target, with a file(GENERATE) cache-busting key covering archs / tunables / AITER pin.

Verdict: Approach is sound. Single-file diff, no behavior change for the prebuilt / AITER_MHA_PATH paths, source-build paths are ordered through a phony target that ck_fused_attn depends on, and install rules stage the freshly built te_libmha_{fwd,bwd}.so (plus the OPTIONAL ck_jit/ runtime tree) from the build tree — that all lines up.

Findings: one minor point posted inline on the CK-JIT file(GLOB) around L189–193 — recommend CONFIGURE_DEPENDS so new JIT scripts get picked up without a manual reconfigure, otherwise the "Ninja rebuilds when inputs change" invariant leaks a little.

Copyright headers: OK — the touched file's AMD line is already 2024-2026.

[ipanfilo]

That is nice improvement. Indeed, a lot of script driven builds run at configure stage. One concern - will moving it to build run different modules/submodules building in parallel? It may be speed improvement but on the other hand can result in build failure due to OOM

[ipanfilo]

Will it be easier to just GLOB all files? Build with CK_JIT does not update repository tree but only temp and install directories

[Micky774]

Updated

[ipanfilo]

Why intermediate location is needed in this scheme, only because we don't know the destitation at the moment of CK_FWD/BWD_SO declaration?

[Micky774]

It's mainly to separate out the build/install tree so that we have an easier time of cleaning the build for incremental work.

[Micky774]

One concern - will moving it to build run different modules/submodules building in parallel? It may be speed improvement but on the other hand can result in build failure due to OOM

Are you referring to the installation of AOTriton? That is indeed a risk since it'll launch in parallel with whatever other aspects of the ninja build can be compiled at that time (so many of the hipcc calls and QoLA building presumably). I'm not sure that the risk is any greater than we already have for our vast hipcc compilation's memory usage though. Or are you referring to something else?