Skip to content

Commit a4e9dd6

Browse files
ipanfiloipanfilo
authored
Update CK-JIT and QoLA (#650)
1 parent 6dd7bcb commit a4e9dd6

7 files changed

Lines changed: 119 additions & 41 deletions

File tree

‎3rdparty/QoLA‎

‎3rdparty/ck_jit‎

‎ci/_utils.sh‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -382,6 +382,20 @@ ck_jit_prebuild() {
382382
if [ "$1" = "build" ]; then
383383
echo "Building CK JIT cache for arch=${_gpu_arch:-<not detected>}..."
384384
python "$_prebuild_py" build --blob-list "$_prebuild_list" $_arch_arg $_jobs_arg > /dev/null
385+
_CK_JIT_CACHE_SNAPSHOT=$(python "$_prebuild_py" cache)
386+
echo "$_CK_JIT_CACHE_SNAPSHOT" | grep Cache
387+
else
388+
if [ -z "${_CK_JIT_CACHE_SNAPSHOT+set}" ]; then
389+
python "$_prebuild_py" cache | grep Cache
390+
else
391+
_cache_now=$(python "$_prebuild_py" cache)
392+
if [ "$_CK_JIT_CACHE_SNAPSHOT" != "$_cache_now" ]; then
393+
echo "Cache diff (build -> now):"
394+
_diff_tmp=$(mktemp)
395+
echo "$_CK_JIT_CACHE_SNAPSHOT" > "$_diff_tmp"
396+
echo "$_cache_now" | diff -u "$_diff_tmp" -
397+
rm -f "$_diff_tmp"
398+
fi
399+
fi
385400
fi
386-
python "$_prebuild_py" cache | grep Cache
387401
}

‎transformer_engine/common/ck_fused_attn/CMakeLists.txt‎

Lines changed: 93 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -22,34 +22,77 @@ if(NOT _gfx1250_idx EQUAL -1)
2222
endif()
2323
set(GPU_TARGETS ${CMAKE_HIP_ARCHITECTURES})
2424
endif()
25-
set(__AITER_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}/../../../3rdparty/QoLA/3rdparty/aiter")
26-
set(__CK_SOURCE_DIR "${__AITER_SOURCE_DIR}/3rdparty/composable_kernel")
2725

28-
set(CK_INCLUDE_DIR "${__CK_SOURCE_DIR}/include")
29-
message(STATUS "ck_include_dir: ${CK_INCLUDE_DIR}")
30-
if(NOT EXISTS "${CK_INCLUDE_DIR}")
31-
message(FATAL_ERROR
32-
"Could not find CK API. "
33-
"Try running 'git submodule update --init --recursive' "
34-
"within the Transformer Engine source.")
26+
set(__QOLA_DIR "${CMAKE_CURRENT_LIST_DIR}/../../../3rdparty/QoLA")
27+
set(__AITER_SOURCE_DIR "${CMAKE_CURRENT_BINARY_DIR}/qola/third_party/aiter")
28+
if (DEFINED ENV{NVTE_AITER_SOURCE_DIR} AND NOT $ENV{NVTE_AITER_SOURCE_DIR} STREQUAL "")
29+
set(__AITER_SOURCE_DIR $ENV{NVTE_AITER_SOURCE_DIR})
30+
message(STATUS "Using AITER source from NVTE_AITER_SOURCE_DIR=${__AITER_SOURCE_DIR}. Disable AITER checkout.")
31+
set(__SKIP_AITER_CHECKOUT TRUE)
32+
else()
33+
set(__SKIP_AITER_CHECKOUT FALSE)
3534
endif()
36-
35+
set(__CK_SOURCE_DIR "${__AITER_SOURCE_DIR}/3rdparty/composable_kernel")
36+
set(CK_INCLUDE_DIR "${__CK_SOURCE_DIR}/include")
3737
set(AITER_INCLUDE_DIR "${__AITER_SOURCE_DIR}/csrc/include")
38-
message(STATUS "aiter_include_dir: ${AITER_INCLUDE_DIR}")
39-
if(NOT EXISTS "${AITER_INCLUDE_DIR}")
40-
message(FATAL_ERROR
41-
"Could not find AITER API. "
42-
"Try running 'git submodule update --init --recursive' "
43-
"within the Transformer Engine source.")
44-
endif()
4538

4639
if(NOT Python_EXECUTABLE)
4740
find_package(Python COMPONENTS Interpreter QUIET)
41+
if (NOT Python_EXECUTABLE)
42+
#Don't fail the entire build if Python isn't found since it may pass if pre-built AITER is used
43+
message(ERROR "Python interpreter not found. Some build steps will be skipped or fail. Please install Python and ensure it is on the PATH.")
44+
endif()
45+
endif()
46+
47+
# Resolve the manifest-pinned AITER commit (defines AITER_SHA) and bring the
48+
# QoLA-managed AITER source tree to that commit before any consumer reads it
49+
# (header validation below, header includes for the .cpp build later, and
50+
# QoLA's own kernel build if the prebuilt cache misses).
51+
set(__QOLA_MANIFEST "${CMAKE_CURRENT_LIST_DIR}/qola_manifest.toml")
52+
set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${__QOLA_MANIFEST}")
53+
file(STRINGS "${__QOLA_MANIFEST}" __AITER_COMMIT_LINES
54+
REGEX "^[ \t]*aiter_commit[ \t]*=[ \t]*\"[^\"]+\"")
55+
list(LENGTH __AITER_COMMIT_LINES __AITER_COMMIT_COUNT)
56+
if(NOT __AITER_COMMIT_COUNT EQUAL 1)
57+
message(FATAL_ERROR
58+
"Expected exactly one 'aiter_commit = \"...\"' line in "
59+
"${__QOLA_MANIFEST}, found ${__AITER_COMMIT_COUNT}.")
60+
endif()
61+
list(GET __AITER_COMMIT_LINES 0 __AITER_COMMIT_LINE)
62+
string(REGEX MATCH "\"([^\"]+)\"" _UNUSED "${__AITER_COMMIT_LINE}")
63+
set(AITER_SHA "${CMAKE_MATCH_1}")
64+
if("${AITER_SHA}" STREQUAL "")
65+
message(FATAL_ERROR
66+
"Failed to parse 'aiter_commit = \"...\"' line in ${__QOLA_MANIFEST}.")
4867
endif()
4968

5069
if(Python_EXECUTABLE)
70+
if (NOT __SKIP_AITER_CHECKOUT)
71+
execute_process(
72+
COMMAND sh -c
73+
"PYTHONPATH=\"${__QOLA_DIR}:$PYTHONPATH\" '${Python_EXECUTABLE}' -m qola.cli checkout \
74+
--manifest '${__QOLA_MANIFEST}' \
75+
--aiter-root '${__AITER_SOURCE_DIR}'"
76+
RESULT_VARIABLE AITER_CHECKOUT_RESULT
77+
OUTPUT_VARIABLE AITER_CHECKOUT_OUTPUT
78+
ERROR_VARIABLE AITER_CHECKOUT_ERROR
79+
OUTPUT_STRIP_TRAILING_WHITESPACE
80+
ERROR_STRIP_TRAILING_WHITESPACE
81+
)
82+
if(NOT AITER_CHECKOUT_RESULT EQUAL 0)
83+
message(FATAL_ERROR
84+
"Failed to sync AITER source tree at ${__AITER_SOURCE_DIR} to "
85+
"manifest-pinned commit ${AITER_SHA}.\n"
86+
"${AITER_CHECKOUT_OUTPUT}\n${AITER_CHECKOUT_ERROR}")
87+
endif()
88+
message(STATUS "[AITER] Synced ${__AITER_SOURCE_DIR} to ${AITER_SHA}")
89+
endif()
90+
5191
execute_process(
52-
COMMAND ${Python_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/check_aiter_mha_args.py --mode both --te-dir "${CMAKE_CURRENT_LIST_DIR}/../../.."
92+
COMMAND ${Python_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/check_aiter_mha_args.py
93+
--mode both
94+
--te-dir "${CMAKE_CURRENT_LIST_DIR}/../../.."
95+
--aiter-root "${__AITER_SOURCE_DIR}"
5396
RESULT_VARIABLE AITER_ARG_CHECK_RESULT
5497
OUTPUT_VARIABLE AITER_ARG_CHECK_OUTPUT
5598
ERROR_VARIABLE AITER_ARG_CHECK_ERROR
@@ -63,8 +106,23 @@ if(Python_EXECUTABLE)
63106
"${AITER_ARG_CHECK_OUTPUT}\n${AITER_ARG_CHECK_ERROR}")
64107
endif()
65108
message(STATUS "AITER API validation passed via check_aiter_mha_args.py")
66-
else()
67-
message(WARNING "Python interpreter not found; skipping AITER API validation.")
109+
endif()
110+
111+
# Sanity-check the resolved include directories now that `qola checkout` has
112+
# materialized the AITER tree.
113+
message(STATUS "ck_include_dir: ${CK_INCLUDE_DIR}")
114+
if(NOT EXISTS "${CK_INCLUDE_DIR}")
115+
message(FATAL_ERROR
116+
"Could not find CK API at ${CK_INCLUDE_DIR}. "
117+
"Re-run the build to let `qola checkout` clone AITER and its "
118+
"composable_kernel submodule.")
119+
endif()
120+
121+
message(STATUS "aiter_include_dir: ${AITER_INCLUDE_DIR}")
122+
if(NOT EXISTS "${AITER_INCLUDE_DIR}")
123+
message(FATAL_ERROR
124+
"Could not find AITER API at ${AITER_INCLUDE_DIR}. "
125+
"Re-run the build to let `qola checkout` clone AITER.")
68126
endif()
69127

70128
set(__AITER_CACHE_DIR "")
@@ -75,12 +133,12 @@ if(NOT "$ENV{NVTE_CK_JIT}" STREQUAL "0")
75133
else()
76134
set(__USE_CK_JIT FALSE)
77135
endif()
136+
78137
if(DEFINED ENV{AITER_MHA_PATH})
79138
message(STATUS "[AITER-BUILD] Using AITER_MHA_PATH=$ENV{AITER_MHA_PATH}")
80139
# use pre-built libraries and includes from a location specified by the user
81140
set(__AITER_CACHE_DIR $ENV{AITER_MHA_PATH})
82-
elseif(NOT __USE_CK_JIT) #disable for CK_JIT for now
83-
# use pre-built cache
141+
elseif(NOT __USE_CK_JIT) # CK_JIT always builds from source; skip the prebuilt cache
84142
include("${CMAKE_CURRENT_LIST_DIR}/aiter_prebuilt.cmake")
85143
get_prebuilt_aiter(__AITER_CACHE_DIR)
86144
elseif(DEFINED ENV{NVTE_AITER_PREBUILT_BASE_URL})
@@ -91,8 +149,6 @@ if(__AITER_CACHE_DIR STREQUAL "")
91149
# If not available or not requested, build from source via QoLA
92150
list(JOIN CMAKE_HIP_ARCHITECTURES ";" GPU_ARCHS_STR)
93151
message(STATUS "[AITER-BUILD] Building AITER kernels for ${GPU_ARCHS_STR} via QoLA.")
94-
set(__QOLA_DIR "${CMAKE_CURRENT_LIST_DIR}/../../../3rdparty/QoLA")
95-
set(__QOLA_MANIFEST "${CMAKE_CURRENT_LIST_DIR}/qola_manifest.toml")
96152
if(__USE_CK_JIT)
97153
message(STATUS "[AITER-BUILD] CK_JIT is enabled; will build CK kernels via CK_JIT.")
98154
set(__CK_JIT_BUILD_DIR "${CMAKE_CURRENT_BINARY_DIR}/ck_jit")
@@ -113,28 +169,30 @@ if(__AITER_CACHE_DIR STREQUAL "")
113169
--tmp-dir "${__CK_JIT_BUILD_DIR}"
114170
--install-dir ${AITER_MHA_INSTALL_DIR}
115171
--jit-name "te_ck_jit"
116-
RESULT_VARIABLE QOLA_BUILD_RESULT
172+
RESULT_VARIABLE CKJIT_BUILD_RESULT
117173
)
174+
if(NOT CKJIT_BUILD_RESULT EQUAL 0)
175+
message(FATAL_ERROR "[AITER-BUILD] CK-JIT build failed.")
176+
endif()
177+
# CK_JIT emits the libs straight into the install dir; nothing to cache.
178+
set(__AITER_MHA_PATH ${AITER_MHA_INSTALL_DIR})
179+
set(__QOLA_INCLUDE_DIR "${__QOLA_BUILD_DIR}/include")
118180
else()
119-
set(__QOLA_BUILD_DIR "${__QOLA_DIR}/build")
181+
set(__QOLA_BUILD_DIR "${CMAKE_CURRENT_BINARY_DIR}/qola")
120182
execute_process(
121183
COMMAND ${CMAKE_COMMAND} -E env "PYTHONPATH=${__QOLA_DIR}:$ENV{PYTHONPATH}"
122184
${Python_EXECUTABLE} -m qola.cli build
123185
--manifest ${__QOLA_MANIFEST}
124186
--aiter-root ${__AITER_SOURCE_DIR}
125187
--output-dir ${__QOLA_BUILD_DIR}
126188
--arch "${GPU_ARCHS_STR}"
189+
--skip-checkout
127190
RESULT_VARIABLE QOLA_BUILD_RESULT
128191
)
129-
endif()
130-
if(NOT QOLA_BUILD_RESULT EQUAL 0)
131-
message(FATAL_ERROR "[AITER-BUILD] QoLA build failed.")
132-
endif()
192+
if(NOT QOLA_BUILD_RESULT EQUAL 0)
193+
message(FATAL_ERROR "[AITER-BUILD] QoLA build failed.")
194+
endif()
133195

134-
if(__USE_CK_JIT)
135-
set(__AITER_MHA_PATH ${AITER_MHA_INSTALL_DIR})
136-
set(__QOLA_INCLUDE_DIR "${__QOLA_BUILD_DIR}/include")
137-
else()
138196
# Copy the final .so libs and exported public headers into the aiter
139197
# prebuilt cache so downstream consumers see a self-contained tree.
140198
get_default_aiter_cache_dir(__AITER_CACHE_DIR)
@@ -147,6 +205,7 @@ if(__AITER_CACHE_DIR STREQUAL "")
147205
set(__QOLA_INCLUDE_DIR "${__AITER_CACHE_DIR}/include")
148206
endif()
149207
else()
208+
message(STATUS "[AITER-BUILD] Using pre-built AITER from ${__AITER_CACHE_DIR}")
150209
set(__AITER_MHA_PATH "${__AITER_CACHE_DIR}/lib")
151210
set(__QOLA_INCLUDE_DIR "${__AITER_CACHE_DIR}/include")
152211
endif()

‎transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,9 @@ string(REGEX MATCH "^[0-9]+\\.[0-9]+" ROCM_VER "${ROCM_VER_CONTENT}")
1919
string(REGEX MATCH "^[0-9]+" ROCM_VER_MAJOR "${ROCM_VER}")
2020

2121
# AITER commit
22-
get_git_commit("${__AITER_SOURCE_DIR}" AITER_SHA)
22+
if(NOT DEFINED AITER_SHA OR "${AITER_SHA}" STREQUAL "")
23+
get_git_commit("${__AITER_SOURCE_DIR}" AITER_SHA)
24+
endif()
2325

2426
# Cache key & local paths
2527
set(AITER_CACHE_ROOT "${CMAKE_CURRENT_LIST_DIR}/../../../build/aiter-prebuilts")

‎transformer_engine/common/ck_fused_attn/check_aiter_mha_args.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ def parse_with_skip_comments(buffer, line, regex, outputs):
3131

3232

3333
def extract_fields_from_header(text: str, struct_name: str) -> List[str]:
34-
struct_field_re = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)\s*(?:=[^;]*)?;\s*$")
34+
struct_field_re = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)\s*(?:=[^;]*|\{[^;]*\})?;\s*$")
3535
struct_end_re = re.compile(r"^\s*};\s*$")
3636

3737
struct_start_re = re.compile(rf"\bstruct\s+{re.escape(struct_name)}\b")
@@ -64,11 +64,14 @@ def main() -> int:
6464
parser = argparse.ArgumentParser(description="Check aiter args usage vs header definition")
6565
parser.add_argument("--mode", choices=["fwd", "bwd", "both"], default="both", help="Mode: fwd, bwd, or both")
6666
parser.add_argument("--te-dir", type=Path, default=Path(__file__).parent.parent.parent.parent, help="Root directory of TransformerEngine")
67+
parser.add_argument("--aiter-root", type=Path, default=None,
68+
help="AITER source tree root. Defaults to <te-dir>/3rdparty/aiter.")
6769
args = parser.parse_args()
70+
aiter_root = args.aiter_root if args.aiter_root else args.te_dir / "3rdparty/aiter"
6871
modes = ["fwd", "bwd"] if args.mode == "both" else [args.mode]
6972
mismatch = 0
7073
for mode in modes:
71-
header_path = args.te_dir / f"3rdparty/QoLA/3rdparty/aiter/csrc/include/mha_{mode}.h"
74+
header_path = aiter_root / f"csrc/include/mha_{mode}.h"
7275
source_path = args.te_dir / f"transformer_engine/common/ck_fused_attn/src/ck_fused_attn_{mode}.cpp"
7376
header_text = header_path.read_text(encoding="utf-8")
7477
source_text = source_path.read_text(encoding="utf-8")

‎transformer_engine/common/ck_fused_attn/qola_manifest.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
[qola]
2-
aiter_commit = "33f2e6af5f39379c739720080ed0033d533f5cb2" # pinned AITER submodule commit
2+
aiter_commit = "d32b0cb62ecc32bb1a858e8437d58eb9b3856af6" # pinned AITER submodule commit
33
namespace = "te"
44
rocm_versions = ["7.2"]
55

0 commit comments

Comments
 (0)