Update QoLA/AITER by Micky774 · Pull Request #599 · ROCm/TransformerEngine

Micky774 · 2026-05-28T16:18:02Z

Description

Updates QoLA as well as moves up the pinned AITER commit

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

ipanfilo · 2026-05-29T04:17:26Z

+# commit QoLA will actually check out and build, not whatever happens to be
+# the submodule's current HEAD at configure time.
+set(__QOLA_MANIFEST "${CMAKE_CURRENT_LIST_DIR}/qola_manifest.toml")
+set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${__QOLA_MANIFEST}")


AITER_SHA is not cached variable. Why is CMAKE_CONFIGURE_DEPENDS needed?

I mainly have it so that direct cmake or ninja building will catch on manifest changes since I use it for incremental builds (and am thinking about including an incremental build option for TE's AITER FA backend). It's not strictly necessary, just a nice to have in such a workflow

Will it also trigger re-executing of QoLA cli in the parent (ck_used_attn/ ) CMAkeLists.txt?

wangye805 · 2026-06-01T16:03:34Z

+      return nullptr;
+    }
+    void* ptr = nullptr;
+    if(hipMallocAsync(&ptr, bytes, stream) != hipSuccess){


Emm, if we let hip runtime to allocate and manage our buffers, this will create a series of issues. For example, if Pytorch or JAX users pre-allocate 97% of HBM, then our hipMallocAsync will return out of memory.

If what new aiter needs for fmha_args.workspace_alloc is a lambda, we can fake it to give jax/pytorch generated workspace buffer ptr?

wangye805 · 2026-06-01T16:04:18Z

+  // callback thread, which holds runtime locks — calling any HIP API from it
+  // (including hipHostFree) deadlocks against concurrent main-thread HIP
+  // calls. Defer the free to ck_tile::pinned_host_releaser's worker thread.
+  fmha_args.pinned_host_alloc = [](size_t bytes) -> std::shared_ptr<void> {


Emm, why do they need host memory allocated? Is it for inference?

Again, can we fake the lambda to use pytorch/jax generated workspace?

VeeraRajasekhar

Verified that ROCm/rocm-libraries#6764 exists in the current pinned aiter commit.

AllenFarcas · 2026-06-16T21:52:57Z

+  size_t head_dim_v,
+  int64_t window_size_left,
+  int64_t window_size_right,
+  bool is_training, bool cuda_graph) {


From my understanding, is_ck_backend_supported is the only consumer of the new graph-capture gate, but it's called without deterministic (TransformerEngine\transformer_engine\common\fused_attn_rocm\fused_attn.cpp line 320) even though nvte_get_fused_attn_backend has it in scope (TransformerEngine\transformer_engine\common\fused_attn_rocm\fused_attn.cpp line 282).

Since deterministic bwd forces the non-graph-safe CK v2 path, is_ck_bwd_graph_capture_safe cannot make the right call without it. See my other related comments above for what could be added to fix.

You should include here the bool deterministic as well

AllenFarcas · 2026-06-16T21:53:02Z

+         is_v3_arch &&
+         dropout == 0.f &&
+         bias_type == NVTE_Bias_Type::NVTE_NO_BIAS &&
+         max_seqlen_q >= 16;


Add && !deterministic

AllenFarcas · 2026-06-16T21:53:06Z

+  // unaffected). Determinism also forces v2 but is invisible here, so it is handled
+  // on the framework side.
+  if(is_training && cuda_graph &&
+     !is_ck_bwd_graph_capture_safe(bias_type, dropout, max_seqlen_q)){


You should include here the deterministic as well

AllenFarcas · 2026-06-16T21:53:12Z

+static bool is_ck_bwd_graph_capture_safe(
+  NVTE_Bias_Type bias_type,
+  float dropout,
+  size_t max_seqlen_q) {


Add bool deterministic

AllenFarcas · 2026-06-16T21:53:15Z

        window_size_left,
-        window_size_right)){
+        window_size_right,
+        is_training, cuda_graph)){


Pass here the deterministic

ipanfilo · 2026-07-03T15:32:13Z

The latest fixes for determinism seem not helping on gfx942

Micky774 · 2026-07-07T14:32:27Z

Alright looks like the failures are unrelated now.

wangye805 · 2026-07-07T18:47:10Z

+    if(bytes == 0){
+      return nullptr;
+    }
+    constexpr size_t kAlign = 256;


Where is this kAlign=256 coming from?

wangye805 · 2026-07-07T18:49:19Z

      is_training, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout,
      h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right,
-      return_max_logit, cuda_graph, false);
+      return_max_logit, false, false);


Do we enable cuda_graph=true before?

It was ignored before, so we can set it to false arbitrarily. I've added a comment clarifying this.

wangye805 · 2026-07-14T15:15:57Z

+// illegal under HIP-graph capture. So reserve a capture-safe upper bound instead.
+//
+// The dq_acc split count is ceil(seqlen_k / kN0) where kN0 is the dispatched bwd KV tile
+// (bn0). The largest bn0 across supported archs is kCkBwdMaxKvTile (per the CK codegen bwd


bn0 is the same as kN0?

Yes, it's a naming difference in different parts of the backend, but it's passed as the same var.

wangye805 · 2026-07-14T15:17:22Z

+// least one extra split for any actually-dispatched kN0 (all <= kCkBwdMaxKvTile), which
+// covers inter-sequence padding of up to a whole tile per sequence -- far above the
+// few-token THD alignment padding used in practice.
+constexpr uint64_t kCkBwdMaxKvTile = 256;


Emm, this const is copied from CK source code? If so, it may be changed later in CK which can cause our issue.

wangye805 · 2026-07-14T15:20:50Z

+  const size_t a16_hdim = (args.d_qk == 192) ? 192 : 128;
+  const size_t dq_acc_seq = args.is_v3_atomic_fp32 ? seqlen_q : a16_seq;
+  const size_t dq_acc_hdim = args.is_v3_atomic_fp32 ? args.d_qk : a16_hdim;
+  const size_t eff_batch = (args.is_group_mode() && args.is_v3_atomic_fp32) ? 1 : args.b;


In my mind, effective batch should be related to whether it's in group mode or not. bf16 or fp32 atomic mode should not affect it.

This is based off of the mha_bwd section: https://github.com/ROCm/aiter/blob/028756633e4192785217838f4924dc16516f5780/csrc/cpp_itfs/mha_bwd.cu#L571-L601

Okay, let's put a comment of this code pointer so in the future we can understand where this piece of hardcoding code is from.

wangye805 · 2026-07-14T15:30:10Z

+  // host-side; v3 (asm) allocates only dq_acc. v3 is tried first but may fall
+  // back to v2, so reserve the larger of the two. The launcher symbol is forced
+  // local by QoLA's export script, so the v2 size is queried through QoLA.
+  const size_t v2_bytes = QOLA_NS(mha_bwd_workspace_size)(make_bwd_traits(args));


Can you remind me where the mha_bwd_workspace_size is defined?

If the v2_bytes sometimes under estimate the real CK needs, can we change the mha_bwd_workspace_size there? Then we don't need to have std::max over three items. (We can just std::max over v3 asm needs and v2 ck needs)

nit: if mha_bwd_workspace_size is defined in CK or qola, maybe should also put v3_dq_acc_bytes into qola

wangye805 · 2026-07-20T20:41:23Z

+  const size_t a16_hdim = (args.d_qk == 192) ? 192 : 128;
+  const size_t dq_acc_seq = args.is_v3_atomic_fp32 ? seqlen_q : a16_seq;
+  const size_t dq_acc_hdim = args.is_v3_atomic_fp32 ? args.d_qk : a16_hdim;
+  const size_t eff_batch = (args.is_group_mode() && args.is_v3_atomic_fp32) ? 1 : args.b;


Okay, let's put a comment of this code pointer so in the future we can understand where this piece of hardcoding code is from.

wangye805 · 2026-07-20T20:42:33Z

+  }
+  const size_t seqlen_q = args.is_group_mode() ? args.max_tokens_q : args.s_q;
+  const size_t elem = args.is_v3_atomic_fp32 ? 4 : 2;
+  const size_t a16_seq = (args.s_q + 15) / 16 * 16;


In the future, let's add non 16 multiple sq's to our unit tests to see what will happen.

wangye805 · 2026-07-20T20:43:35Z

+  // host-side; v3 (asm) allocates only dq_acc. v3 is tried first but may fall
+  // back to v2, so reserve the larger of the two. The launcher symbol is forced
+  // local by QoLA's export script, so the v2 size is queried through QoLA.
+  const size_t v2_bytes = QOLA_NS(mha_bwd_workspace_size)(make_bwd_traits(args));


nit: if mha_bwd_workspace_size is defined in CK or qola, maybe should also put v3_dq_acc_bytes into qola

wangye805 · 2026-07-20T20:44:33Z

  fmha_args.mask_type = static_cast<int>(static_cast<mask_enum>(args.attn_mask_type));
-  fmha_args.use_asm_v3 = args.uses_bwd_v3;
+  // Mirrors AITER's small-seqlen guard at aiter/ops/mha.py:1689.
+  fmha_args.use_asm_v3 = (args.s_q < 16) ? false : args.uses_bwd_v3;


nit: have such a testcase (s_q<16) to see if things work as our expectation

We already have several where s_q=1 which trigger this code path. Specifically it was put in as a response to a failing test case originally.

Micky774 added 11 commits April 28, 2026 17:00

Updated QoLA (to port CK receipt patch) and TE manifest

d89f90f

Updated manifest

312212f

Corrected AITER mha args validation against pinned commit

89f6983

Updated cmake w/ dubious ownership protection

c417e40

Merge branch 'dev' into zain/qola/aiter-update

90e3a5a

Corrected logging

c7ecaf7

Updated qola to build aiter w/ new third_party spec

da6e9a6

Added guards against AITER known buggy implementations

e7ed124

Updated build

6241f99

Drop guard for corrected bug

82ebdb6

Merge branch 'dev' into zain/qola/aiter-update

68a2eaf

Micky774 requested review from ipanfilo, wangye805 and wenchenvincent as code owners May 28, 2026 16:18

Micky774 added the ci-level 3 CI test level 3 label May 28, 2026

Update AITER commit, adopt new API

f9ab59c

ipanfilo reviewed May 29, 2026

View reviewed changes

Micky774 requested review from AllenFarcas and VeeraRajasekhar May 29, 2026 18:26

wangye805 requested changes Jun 1, 2026

View reviewed changes

Micky774 added 2 commits June 5, 2026 19:21

Added guard against graph-unsafe CK V2 kernels

fec80b8

Merge branch 'dev' into zain/qola/aiter-update

6eb3b23

VeeraRajasekhar reviewed Jun 9, 2026

View reviewed changes

Micky774 added 2 commits June 16, 2026 16:20

Updated with AOT memory handling (dynamic on host)

8fd79f9

Merge branch 'dev' into zain/qola/aiter-update

ab68ca9

AllenFarcas reviewed Jun 16, 2026

View reviewed changes

Moved workspace zeroing from TE to QoLA CK patch

2411723

Micky774 added 2 commits July 6, 2026 14:51

Host memory zero on TE side

9860a92

Reduce bound split size

6a23ce3

wangye805 requested changes Jul 7, 2026

View reviewed changes

Micky774 added 9 commits July 7, 2026 20:54

Corrected memory allocation upper bound

52b39d5

Added ignored comment to ignored cuda_graph arg

9c8fb58

Updated to leverage CK patch for device-side workspace prep

729dfe2

Merge branch 'dev' into zain/qola/aiter-update

270c358

Zeroed out dk_expanded buffer, and enforced stricter aiter zeroing

49e65f5

Added back dropped seqlen correction

db13f73

Remove unnecessary zeroing

0f7f531

Update qola patch to CK dq_acc zeroing

4fe3a7b

Lower bound sizing fix

cf46085

Micky774 force-pushed the zain/qola/aiter-update branch from fec6d6b to cf46085 Compare July 13, 2026 21:10

wangye805 requested changes Jul 14, 2026

View reviewed changes

Micky774 added 4 commits July 17, 2026 17:16

Corrected JAX group mode sizing issue

9c2cccb

Corrected zeroing via AITER patch

e3df798

Merge branch 'dev' into zain/qola/aiter-update

314777b

Corrected patch

3346d3c

wangye805 approved these changes Jul 20, 2026

View reviewed changes

ipanfilo reviewed Jul 21, 2026

View reviewed changes

Comment thread 3rdparty/ck_jit

Comment thread transformer_engine/jax/csrc/extensions/attention.cpp Outdated

Comment thread ci/ck_jit_prebuild.txt

Micky774 mentioned this pull request Jul 21, 2026

Updated to use new device workspace prep function ROCm/ck-jit#7

Closed

1 task

Micky774 added 4 commits July 22, 2026 16:53

Kept default nullptr for nvidia builds

c7b8e1d

Updated prebuilt list for gfx9 kernels

e2f8832

Updated ck-jit to main

0e40940

Merge branch 'dev' into zain/qola/aiter-update

d65de00

Micky774 requested a review from ipanfilo July 22, 2026 20:56

ipanfilo approved these changes Jul 22, 2026

View reviewed changes

Uh oh!

Conversation

Micky774 commented May 28, 2026

Description

Type of change

Changes

Checklist:

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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

VeeraRajasekhar left a comment

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ipanfilo commented Jul 3, 2026

Uh oh!

Micky774 commented Jul 7, 2026

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!

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!

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