Update AITER subcommit and refactor internal AITER/CK FA API usage by Micky774 · Pull Request #446 · ROCm/TransformerEngine

Micky774 · 2026-02-09T21:27:13Z

Description

Updates AITER subcommit as well as refactors our internal usage of the API for greater clarity and explicitness. TODO: Update AITER commit upon merger of ROCm/aiter#2055

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:

Updates AITER subcommit
Refactors to central implementation functions for batch vs group mode fwd/bwd functions
Refactors to explicit instantiation of arguments (rather than relying on positional)
Includes new AITER mask type conversion as part of necessary arguments for the new API
Streamlines logging
Adds helper python script to explicitly compare our application of the AITER API vs the AITER API's declaration through its header files.
Uses helper script during install as sanity-check on appropriate usage of AITER API
WIP pending tests

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

Micky774

A few comments to hopefully help reviewers.

transformer_engine/__init__.py

Micky774 · 2026-02-09T21:31:08Z

tools/check_aiter_mha_args_usage.py

This helper script scans semi-hard-coded files wrt TE source-code in order to directly compare AITER's internal API and our attempt at utilizing it. This script is run during setup through setup.py

You can put this in the comment of this file.

Also don't forget to add the copyright

Micky774 · 2026-02-09T21:35:07Z

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

+int ck_to_aiter_mask_type(mask_enum mask_type, ck_tile::index_t left, ck_tile::index_t right){
+  if(
+    mask_type == mask_enum::no_mask ||
+    mask_type == mask_enum::window_generic
+  ) return 0;
+  if(left == -1 && right == 0){
+    return mask_type == mask_enum::mask_top_left ? 1 : 2;
+  }
+  return 3;
+}


This is based on their op_tests/cpp/mha/* which isn't the most stable or reliable source, but it's the only concrete location where such a mapping is used or documented.

Micky774 · 2026-02-09T21:35:56Z

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

+void log_bwd_config(const char* func_name, const aiter::mha_bwd_args& fmha_args, bool ck_log_config){
+  if (!ck_log_config) {
+    return;
  }

+  auto log_value = [](const char* label, const auto& value) {
+    std::cout << label << ": " << value << "\n";
+  };


This merely standardizes the logging to make it a bit easier to parse through it at a glance while guaranteeing uniformity. This is similarly implemented across both files. Note the signature has been reduced to account for only what isn't stored in fmha_args.

Same here. Put those helpful description into the comments

Micky774 · 2026-02-09T21:37:41Z

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

@@ -499,10 +508,13 @@ hipError_t ck_attn_bwd(
  bool uses_bwd_v3,
  bool is_v3_atomic_fp32,
  int how_v3_bf16_cvt,
+  bool is_group_mode,
+  const char* func_name,
+  bool ck_log_config,
  hipStream_t stream){


This abstracted implementation function has a signature that is a superset of the other higher-level API functions in the file, so that both can route directly to this function without affecting the API outside of this file.

Micky774 · 2026-02-09T21:45:29Z

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

-                         std::pair<const void*, const void*>{philox_seed_ptr, philox_offset_ptr}};
-  }();
+  aiter::mha_bwd_args fmha_args{};
+  fmha_args.mask_type = ck_to_aiter_mask_type(mask_type, left, right);


This is the AITER mask type, despite the same argument referring to the CK mask type in the FWD pass

Note that there is a PR for removing this argument which is still pending.

Micky774 · 2026-02-09T21:47:38Z

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

  }
+  return hipSuccess;
+}
+hipError_t ck_attn_bwd(


Both ck_attn_{varlen_}bwd are wrappers around _ck_attn_bwd_impl + TE-side post-processing kernels. Those kernels can probably be refactored similarly but that's a bit outside the scope here and of dubious value.

setup.py

Micky774 · 2026-02-11T16:02:12Z

setup.py

+            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)


Explicitly checks the AITER API usage

Actually you can put your comment in PR to the comments in the source codes

setup.py

tools/check_aiter_mha_args_usage.py

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp

wangye805 · 2026-02-17T03:57:48Z

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_utils.cpp

+    dladdr((void*)set_aiter_asm_dir, &info);
+    auto install_lib_path = std::filesystem::path(info.dli_fname).parent_path() / "aiter";
+    const char* log_ck_config = std::getenv("NVTE_LOG_CK_CONFIG");
+    auto editable_install_path = std::filesystem::path(info.dli_fname).parent_path().parent_path().parent_path() / "3rdparty" / "aiter" / "hsa";


I recall we also copy hsa in editable installation

I don't think that we do

I don't have a dev TE sandbox at hand but here is one with v2.6 release editable installation:

Maybe there's a difference in building from source vs downloading the pre-built AITER lib? My current editable install doesn't have the copied aiter directory under transformer_engine/lib.

Okay, I think now I understand the root cause. Your aiter commit update need to bring back the hsa dir copy which was removed in Dragan's PR: https://github.com/ROCm/TransformerEngine/pull/402/changes#diff-e0641ac80f912d730377ecbf45edc9573b4ad4a07ee875e474a49f3d57605f1f

Otherwise, the non-editable installation will not have hsa dir in transformer_engine installation dir either

With the hsa dir copy functionality brought back, we don't need to distinguish the install and editable install path, I believe?

wangye805 · 2026-02-17T04:00:05Z

transformer_engine/common/ck_fused_attn/aiter_build.sh


 rm -rf "${AITER_DIR}/aiter/jit/build"
-AITER_LOG_MORE=1 \
+AITER_LOG_MORE=0 \


Is it possible to link pip install -v to AITER_LOG_MORE?

wangye805 · 2026-02-17T04:01:01Z

setup.py

+            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)


Actually you can put your comment in PR to the comments in the source codes

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp

wangye805 · 2026-02-18T05:21:42Z

tools/check_aiter_mha_args_usage.py

You can put this in the comment of this file.

Also don't forget to add the copyright

wangye805 · 2026-02-18T05:22:35Z

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

+void log_bwd_config(const char* func_name, const aiter::mha_bwd_args& fmha_args, bool ck_log_config){
+  if (!ck_log_config) {
+    return;
  }

+  auto log_value = [](const char* label, const auto& value) {
+    std::cout << label << ": " << value << "\n";
+  };


Same here. Put those helpful description into the comments

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

wangye805 · 2026-02-18T05:32:14Z

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

+  if (!is_group_mode) {
+    std::tie(bias_type, bias_shape) = get_ck_bias_type_shape(attn_bias_type, b, h, bias_b, bias_h);
+  }
+  const ck_tile::index_t batch_stride_bias = [&]() -> ck_tile::index_t {


This overlaps with the lambda defined in fwd:

TransformerEngine/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp

Lines 187 to 217 in 29878cf

const ck_tile::index_t batch_stride_bias = [&]() -> ck_tile::index_t {

if(is_group_mode){

return 0;

}

switch (bias_shape) {

case BiasShape::k11SS:

case BiasShape::k1HSS:

return 0;

case BiasShape::kBHSS:

return bias_h * max_seqlen_q * max_seqlen_k;

case BiasShape::kB1SS:

return max_seqlen_q * max_seqlen_k;

default:

throw std::runtime_error("Invalid bias shape");

}

}();

const ck_tile::index_t nhead_stride_bias = [&]() -> ck_tile::index_t {

if(is_group_mode){

return 0;

}

switch (bias_shape) {

case BiasShape::k1HSS:

case BiasShape::kBHSS:

return max_seqlen_q * max_seqlen_k;

case BiasShape::k11SS:

case BiasShape::kB1SS:

return 0;

default:

throw std::runtime_error("Invalid bias shape");

}

}();

Try to consolidate them into util.cpp?

Also it's better to avoid using lambda with catching every reference. Previously I used it because I was copying from ck example.

wangye805 · 2026-02-18T05:39:27Z

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp

+  bias_enum bias_type = bias_enum::no_bias;
+  BiasShape bias_shape = BiasShape::k11SS;
+  if(is_group_mode){
+    seqstart_q_ptr = cu_seqlen_q_padded_ptr==nullptr? cu_seqlen_q_ptr: cu_seqlen_q_padded_ptr;


Make it consistent with bwd:

TransformerEngine/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

Lines 606 to 620 in 29878cf

if (is_group_mode) {

fmha_args.seqstart_q_ptr = cu_seqlen_q_padded_ptr==nullptr? cu_seqlen_q_ptr: cu_seqlen_q_padded_ptr;

fmha_args.seqstart_k_ptr = cu_seqlen_kv_padded_ptr==nullptr? cu_seqlen_kv_ptr: cu_seqlen_kv_padded_ptr;

fmha_args.seqlen_q_ptr = nullptr;

fmha_args.seqlen_k_ptr = nullptr;

fmha_args.cu_seqlen_q_ptr = cu_seqlen_q_ptr;

fmha_args.cu_seqlen_k_ptr = cu_seqlen_kv_ptr;

} else {

fmha_args.seqstart_q_ptr = nullptr;

fmha_args.seqstart_k_ptr = nullptr;

fmha_args.seqlen_q_ptr = nullptr;

fmha_args.seqlen_k_ptr = nullptr;

fmha_args.cu_seqlen_q_ptr = nullptr;

fmha_args.cu_seqlen_k_ptr = nullptr;

}

wangye805 · 2026-02-18T05:42:58Z

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

-    const ck_tile::index_t stride_q = stride_s_q;
-    const ck_tile::index_t stride_k = stride_s_k;
-    const ck_tile::index_t stride_v = stride_s_v;
-    // bias of shape (bias_b, bias_h, s_q, s_kv)


Can you keep this kind of comments so that later we can understand why we set stride_bias/dbias to max_seqlen_k

wangye805 · 2026-02-18T05:43:28Z

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

-    const ck_tile::index_t stride_dv_expanded = stride_s_dv_expanded;
-    const ck_tile::index_t stride_dq_acc = d_qk; //dq_acc of shape (nsplits, B, H, S, D)
-    // dbias is of the same shape as bias
-    // but ck only take dbias with BHSS


Also this comment is crucial for dbias reduction

wangye805 · 2026-02-18T05:51:32Z

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp

+    nullptr, nullptr,
+    nullptr, nullptr,


Add some comments so code readers can know quickly what's those 4 nullptrs are for?

wangye805 · 2026-02-18T05:54:13Z

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_utils.cpp

+    dladdr((void*)set_aiter_asm_dir, &info);
+    auto install_lib_path = std::filesystem::path(info.dli_fname).parent_path() / "aiter";
+    const char* log_ck_config = std::getenv("NVTE_LOG_CK_CONFIG");
+    auto editable_install_path = std::filesystem::path(info.dli_fname).parent_path().parent_path().parent_path() / "3rdparty" / "aiter" / "hsa";


I don't have a dev TE sandbox at hand but here is one with v2.6 release editable installation:

tools/check_aiter_mha_args_usage.py

setup.py

transformer_engine/common/gemm/ck_grouped_gemm.cpp

transformer_engine/common/ck_fused_attn/CMakeLists.txt

Micky774 · 2026-03-06T20:45:21Z

level 3 run for tracking: https://github.com/ROCm/TransformerEngine/actions/runs/22781278349

transformer_engine/common/gemm/ck_grouped_gemm.cpp

Micky774 · 2026-03-17T19:10:45Z

Level 3 run for tracking: https://github.com/ROCm/TransformerEngine/actions/runs/23265047818 (updated for commit 638c9e6

Micky774 · 2026-03-20T18:17:17Z

Okay so after patching the race condition w/ Xinya's mutex solution, another bug was revealed through the JAX MGPU tests which is actually a hip launch error due to device ordinal mismatch. This is essentially because the kernel map is keyed by the kernel name only, NOT by device ID, meaning that multiple device threads will map different kernels to the same key, and when loading we end up with the wrong kernels on the wrong devices. I've added a patch to address this issue, with a corresponding AITER PR here: ROCm/aiter#2401

Micky774 added 10 commits February 2, 2026 15:46

Initial inclusion of new API in fwd as well as part 1 of refactor

f02c48e

Initial implementation of refactor/API update across ALL CK funcs

0b0ad93

Updated logging

c198cbd

Add script for comparing AITER/TE API

a52bb32

Reconcile new AITER mask type

77f0a05

Updated API helper tool

1637266

Merge branch 'dev' into zain/aiter-api

568e9b5

Formatting

2cb6d82

Added sys exit to script

cf4aa9e

Slightly better error message

e25cea8

Micky774 commented Feb 9, 2026

View reviewed changes

Updated AITER_ASM_DIR implementation

2122479

Micky774 marked this pull request as ready for review February 11, 2026 16:01

Micky774 requested review from ipanfilo, wangye805 and wenchenvincent as code owners February 11, 2026 16:01

Micky774 commented Feb 11, 2026

View reviewed changes

Update AITER

4817e72

Micky774 changed the title ~~[WIP] Update AITER subcommit and refactor internal AITER/CK FA API usage~~ [No Merge] Update AITER subcommit and refactor internal AITER/CK FA API usage Feb 11, 2026

Micky774 added 3 commits February 12, 2026 10:26

Updated AITER_ASM_DIR logic to allow for hip-free use

837b827

Re-introduce setup AITER API check

68ca0fe

Update AITER to custom feature branch

ae688ab

Micky774 changed the title ~~[No Merge] Update AITER subcommit and refactor internal AITER/CK FA API usage~~ Update AITER subcommit and refactor internal AITER/CK FA API usage Feb 16, 2026

Micky774 added 2 commits February 16, 2026 16:06

Reduce AITER build verbosity

762b91b

Updated API

0a7187d

wangye805 requested changes Feb 17, 2026

View reviewed changes

Micky774 added 2 commits February 17, 2026 11:22

Address PR comments

39b27bc

Updated bias stride calculations

29878cf

wangye805 requested changes Feb 18, 2026

View reviewed changes

Merge branch 'dev' into zain/aiter-api

6846a27

ipanfilo reviewed Feb 20, 2026

View reviewed changes

tools/check_aiter_mha_args_usage.py Show resolved Hide resolved

VeeraRajasekhar self-requested a review February 23, 2026 06:45

Updated ck_fused_attn lib build to include copying HSA

78f1d69

Micky774 force-pushed the zain/aiter-api branch from ab79b40 to 78f1d69 Compare February 27, 2026 22:34

Micky774 added 9 commits March 3, 2026 15:24

Corrected AITER bug and moved to TE feature branch

f935956

Merge branch 'dev' into zain/aiter-api

b90da33

Added back dropped code from merge conflict

0475f85

Downgrade to more conservative AITER commit for compat

5db08ea

Removed python-level args check

d5e5ec6

Removed old tools

151e9ca

Corrected arg_size types manually

5a18d16

Updated AITER commit and fixed API mismatch in group gemm

db34177

Added build-time AITER API usage check

9cd2833

ipanfilo reviewed Mar 5, 2026

View reviewed changes

Micky774 added 4 commits March 5, 2026 14:16

PR review comments

a6b831e

Undo extra import removal

d1bd569

Adjusted python requirement in cmakelist

686e6a2

Updated group-gemm dispatch

52e1a1a

ipanfilo reviewed Mar 6, 2026

View reviewed changes

transformer_engine/common/ck_fused_attn/CMakeLists.txt Outdated Show resolved Hide resolved

Made AITER API check earlier

48c5839

matthiasdiener reviewed Mar 6, 2026

View reviewed changes

transformer_engine/common/gemm/ck_grouped_gemm.cpp Outdated Show resolved Hide resolved

Micky774 mentioned this pull request Mar 10, 2026

Utilize AITER HSA embeddings #477

Open

13 tasks

Micky774 added 2 commits March 16, 2026 15:05

Update AITER w/ Xinya's patch

9b2166e

Merge branch 'dev' into zain/aiter-api

3dacb2a

wangye805 approved these changes Mar 17, 2026

View reviewed changes

Micky774 added 2 commits March 18, 2026 15:16

Revert and cherry-pick aiter subcommit

638c9e6

Patched unordered map mgpu kernel collision

c786c3b

	const ck_tile::index_t batch_stride_bias = [&]() -> ck_tile::index_t {
	if(is_group_mode){
	return 0;
	}
	switch (bias_shape) {
	case BiasShape::k11SS:
	case BiasShape::k1HSS:
	return 0;
	case BiasShape::kBHSS:
	return bias_h * max_seqlen_q * max_seqlen_k;
	case BiasShape::kB1SS:
	return max_seqlen_q * max_seqlen_k;
	default:
	throw std::runtime_error("Invalid bias shape");
	}
	}();
	const ck_tile::index_t nhead_stride_bias = [&]() -> ck_tile::index_t {
	if(is_group_mode){
	return 0;
	}
	switch (bias_shape) {
	case BiasShape::k1HSS:
	case BiasShape::kBHSS:
	return max_seqlen_q * max_seqlen_k;
	case BiasShape::k11SS:
	case BiasShape::kB1SS:
	return 0;
	default:
	throw std::runtime_error("Invalid bias shape");
	}
	}();

	if (is_group_mode) {
	fmha_args.seqstart_q_ptr = cu_seqlen_q_padded_ptr==nullptr? cu_seqlen_q_ptr: cu_seqlen_q_padded_ptr;
	fmha_args.seqstart_k_ptr = cu_seqlen_kv_padded_ptr==nullptr? cu_seqlen_kv_ptr: cu_seqlen_kv_padded_ptr;
	fmha_args.seqlen_q_ptr = nullptr;
	fmha_args.seqlen_k_ptr = nullptr;
	fmha_args.cu_seqlen_q_ptr = cu_seqlen_q_ptr;
	fmha_args.cu_seqlen_k_ptr = cu_seqlen_kv_ptr;
	} else {
	fmha_args.seqstart_q_ptr = nullptr;
	fmha_args.seqstart_k_ptr = nullptr;
	fmha_args.seqlen_q_ptr = nullptr;
	fmha_args.seqlen_k_ptr = nullptr;
	fmha_args.cu_seqlen_q_ptr = nullptr;
	fmha_args.cu_seqlen_k_ptr = nullptr;
	}

Conversation

Micky774 commented Feb 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

Micky774 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Micky774 commented Feb 9, 2026 •

edited

Loading

Micky774 commented Mar 17, 2026 •

edited

Loading