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-29T03:50:06Z

+    COMMAND sh -c
+      "tmp=$(mktemp /tmp/gitconfig.XXXXXX) || exit 1; \
+GIT_CONFIG_GLOBAL=$tmp git config --global --add safe.directory '*' >/dev/null 2>&1; \
+GIT_CONFIG_GLOBAL=$tmp PYTHONPATH=\"${__QOLA_DIR}:$PYTHONPATH\" '${Python_EXECUTABLE}' -m qola.cli checkout \


Wouldn't it make sense to integrate safe.directory overriding to qola? The pattern with dubious ownership is probably not TE specific

I worry that such behavior is a bit too authoritative for qola if that makes sense? My reasoning is that the permission scope here seems to be outside of qola and hence qola should not be the one in charge of it. I'm open to reconsidering that, but it's just my initial position.

It is valid concern bearing in mind QoLA is intended to be reused by different components. May be make this behavior controllable then. It is OK to keep the things in TE, it will just require doing things this overriding twice - here and when build is called - BTW, there are comments there but not actual code change

Sounds good, but can we do that in a separate PR to QoLA first?

Now, with qola expected to create aiter dir runtime (it is not a part of repository tree), is there a reason to modify git config safe.directory? It was added for specific case: repository tree is checked out by a user and then shared with a container run with another user credentials.

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

AllenFarcas · 2026-06-16T22:00:16Z

+    # graph capture this corner is invisible to the C++ backend selector (determinism
+    # is not part of its signature), so fall back to the unfused backend here.


Similar comment here

AllenFarcas · 2026-06-22T16:25:17Z

+  fmha_args.d_sink_ptr = nullptr;
  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.


The next time AITER version is bumped this hardcoded line number will not be relevant anymore.

ipanfilo · 2026-06-24T01:42:31Z

    void* ptr = nullptr;
    if(hipHostMalloc(&ptr, bytes, hipHostMallocDefault) != hipSuccess){
-      throw std::runtime_error("ck_fused_attn bwd: hipHostMalloc failed for AITER pinned host buffer.");
+      throw std::runtime_error("ck_fused_attn bwd: pinned host staging was not reserved and "


This message is specific for hipErrorStreamCaptureUnsupported only while there might be other allocation errors, so dispatch by hipGetLastError() is needed

I've updated the error string to be more descriptive, lmk if this is sufficient.

ipanfilo · 2026-06-24T01:43:50Z

    bool resolves_to_v3 = fmha_args.use_asm_v3 && !fmha_args.is_deterministic &&
                          !fmha_args.has_dbias && fmha_args.bias_type == 0 &&
                          !fmha_args.has_dropout && is_v3_arch;
-    if(!resolves_to_v3){


W/o this throw the whole graph capture safety net is useless

The graph capture safety net actually needs to be removed entirely since now the CK kernels provided should be graph safe. I've updated the PR to remove it.

# Conflicts: # 3rdparty/QoLA # transformer_engine/common/ck_fused_attn/CMakeLists.txt # transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake # transformer_engine/common/ck_fused_attn/check_aiter_mha_args.py # transformer_engine/common/ck_fused_attn/qola_manifest.toml

ipanfilo · 2026-06-27T00:17:38Z

+# QoLA-managed AITER source tree to that commit before any consumer reads it
+# (header validation below, header includes for the .cpp build later, and
+# QoLA's own kernel build if the prebuilt cache misses).
+include("${CMAKE_CURRENT_LIST_DIR}/aiter_prebuilt.cmake")


It is not needed here - AITER_SHA is determined above, cache directory is only needed if neither AITER_MHA_PATH is set nor CK-JIT is enabled, so it is included at line 158

ipanfilo · 2026-06-27T00:21:46Z

+    COMMAND sh -c
+      "tmp=$(mktemp /tmp/gitconfig.XXXXXX) || exit 1; \
+GIT_CONFIG_GLOBAL=$tmp git config --global --add safe.directory '*' >/dev/null 2>&1; \
+GIT_CONFIG_GLOBAL=$tmp PYTHONPATH=\"${__QOLA_DIR}:$PYTHONPATH\" '${Python_EXECUTABLE}' -m qola.cli checkout \


Now, with qola expected to create aiter dir runtime (it is not a part of repository tree), is there a reason to modify git config safe.directory? It was added for specific case: repository tree is checked out by a user and then shared with a container run with another user credentials.

When CK-JIT cannot find workspace-query symbols (dq_ws_host_size, needs_zero_dq_acc) for a kernel variant — which happens for deterministic kernels when d_v differs from d_qk (e.g. d_qk=128, d_v=64) — AITER's mha_bwd_workspace_size() reports only the non-deterministic workspace size (~132 MiB), causing the kernel to overflow into adjacent XLA buffers and corrupt the forward output (NaN in value_and_grad). Fix: add the old explicit dq_acc formula as a floor in ck_attn_bwd_workspace_size(), matching the per-sequence kv split count used by the CK v2 deterministic accumulation algorithm. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

ipanfilo

It seems there are still some numeric failures in determenistic path

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

The latest fixes for determinism seem not helping on gfx942

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

Micky774 added 5 commits June 18, 2026 17:33

Added entry-point for workspace calc in QoLA

f521ef0

Updated AITER commit to one w/ CK graph fix cherry-pick

e5ebe9c

Updated determnistic arg plumbing

c08548e

Updated CK JIT

292870e

Updated to gfx-1250 qola branch

3476f53

AllenFarcas reviewed Jun 22, 2026

View reviewed changes

Added pinned host memory pool allocation scheme

79c400d

ipanfilo reviewed Jun 24, 2026

View reviewed changes

ipanfilo and others added 2 commits June 24, 2026 08:55

Update CK JIT and prebuiuld cache list

3a98eda

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

bb0797a

Micky774 force-pushed the zain/qola/aiter-update branch from f004fcf to bb0797a Compare June 24, 2026 15:23

Micky774 added 3 commits June 26, 2026 16:57

Made error more descriptive

bf2ae3c

Removed outdated graph safety guards

8a91a58

Marked unused param in get attn backend

679283a

Micky774 force-pushed the zain/qola/aiter-update branch from a8c5290 to 679283a Compare June 26, 2026 17:17

Micky774 added 2 commits June 26, 2026 18:52

Updated qola commit

eaa860d

Micky774 requested a review from ipanfilo June 26, 2026 19:56

ipanfilo reviewed Jun 27, 2026

View reviewed changes

ipanfilo and others added 3 commits June 29, 2026 09:12

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

48e73b2

PR feedback

9172b4b

ipanfilo reviewed Jun 30, 2026

View reviewed changes

Comment thread transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp Outdated

Micky774 added 3 commits July 1, 2026 19:22

Fixed flaky determinism failures for padded configs

fe94903

Corrected comment

eaf3cce

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

		# graph capture this corner is invisible to the C++ backend selector (determinism
		# is not part of its signature), so fall back to the unfused backend here.

Uh oh!

Conversation

Micky774 commented May 28, 2026

Description

Type of change

Changes

Checklist:

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!

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!

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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ipanfilo left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ipanfilo commented Jul 3, 2026

Uh oh!