Utilize AITER HSA embeddings by Micky774 · Pull Request #477 · ROCm/TransformerEngine

Micky774 · 2026-03-09T19:21:39Z

Description

Currently AITER relies on reading AITER_ASM_DIR which is not in general thread-safe when we (and potentially others) modify the variable at runtime. AITER supports embedding the HSA table into a header which will be used directly instead of relying on runtime loading of the kernels via AITER_ASM_DIR. This results in thread-safe operation, without significant cost.

There is an AITER PR open for compatibility on their side, however that would still require us to use newer versions of AITER which is pending this PR.

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:

Uses PyTorch's generate_aiter_embedded_hsa.py to generate header file for the HSA embedding
Updates build system to use HSA generating script, and properly pass to compile.py
Requires the use flag passthrough in AITER's compile.py which is pending a PR.
AITER diff for reference.

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

wangye805

Okay, after seeing ROCm/aiter@a64fa18...1b616b0#diff-fbc2a753d0cd5161aad3c68beb6eab5a577bae39fe4cb7b1af1c3b6c8ae963beR72, I sort of understand that aiter compile_ops provide hooks to modify the source generation, compilation flags, and so on in their json file.

In this case, how about move the archive (static lib) part into their compile.py? We can provide another option or env to their compile.py to go the static lib path. Therefore TE side is just loading their lib.a.

transformer_engine/common/ck_fused_attn/generate_aiter_embedded_hsa.py

wangye805 · 2026-03-10T13:59:35Z

transformer_engine/common/ck_fused_attn/CMakeLists.txt

+             AITER_EMBEDDED_HSA_HEADER=aiter_embedded_hsa.h
+             AITER_EMBEDDED_HSA_INCLUDE_DIR=${AITER_EMBEDDED_HSA_HEADER_DIR})
+      else()
+        message(STATUS "[AITER-PREBUILT] No supported V3 ASM arch selected; skipping embedded HSA generation.")


[AITER-PREBUILT]? [AITER-BUILD]

transformer_engine/common/ck_fused_attn/generate_aiter_embedded_hsa.py

transformer_engine/common/ck_fused_attn/CMakeLists.txt

Micky774 · 2026-03-10T20:10:27Z

Okay, after seeing ROCm/aiter@a64fa18...1b616b0#diff-fbc2a753d0cd5161aad3c68beb6eab5a577bae39fe4cb7b1af1c3b6c8ae963beR72, I sort of understand that aiter compile_ops provide hooks to modify the source generation, compilation flags, and so on in their json file.

In this case, how about move the archive (static lib) part into their compile.py? We can provide another option or env to their compile.py to go the static lib path. Therefore TE side is just loading their lib.a.

To do that, we'd need to essentially implement this PR on the AITER side anyways. We can't use their existing hooks exposed in compile.py to actually build static archives. The "hack" on our end is still the simplest way of accomplishing it.

ipanfilo · 2026-03-10T19:33:56Z

3rdparty/aiter

You probably don't need path splitting there either, it should work with DAITER_EMBEDDED_HSA_HEADER=\"{header_path_var}\"'
Also, I wonder if it can be overridden w/o AITERT scripts change with using HIPCC_COMPILE_FLAGS_APPEND or whatever

ipanfilo · 2026-03-10T19:41:10Z

transformer_engine/common/ck_fused_attn/CMakeLists.txt

-                  --output ${AITER_EMBEDDED_HSA_HEADER}
+                  --output ${AITER_EMBEDDED_HSA_HEADER_PATH}
                  --subdirs ${AITER_EMBEDDED_HSA_SUBDIRS}
          RESULT_VARIABLE AITER_MAKE_HSA_RET


Nit: You can use COMMAND_ERROR_IS_FATAL ANY instead so CMake will fail itself on command failure

wangye805 · 2026-03-10T20:57:24Z

Okay, after seeing ROCm/aiter@a64fa18...1b616b0#diff-fbc2a753d0cd5161aad3c68beb6eab5a577bae39fe4cb7b1af1c3b6c8ae963beR72, I sort of understand that aiter compile_ops provide hooks to modify the source generation, compilation flags, and so on in their json file.
In this case, how about move the archive (static lib) part into their compile.py? We can provide another option or env to their compile.py to go the static lib path. Therefore TE side is just loading their lib.a.

To do that, we'd need to essentially implement this PR on the AITER side anyways. We can't use their existing hooks exposed in compile.py to actually build static archives. The "hack" on our end is still the simplest way of accomplishing it.

I'm not saying just building the static library in aiter optest compile.py. I'm referring to adding the archive part after lib.so's are finished building. Perhaps after https://github.com/ROCm/aiter/blob/9c2fef901dfbc1cc0e212038173d9d46160e3e12/op_tests/cpp/mha/compile.py#L98, then run the archive to get the lib.a

wangye805 · 2026-03-10T21:00:14Z

transformer_engine/common/ck_fused_attn/generate_aiter_embedded_hsa.py

@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.


Add "this file was modified for rocm Transformer Engine", like

TransformerEngine/setup.py

Line 1 in 9b31283

# This file was modified for portability to AMDGPU

Ilya had more experience on the copyright area. Can you comment on this @ipanfilo ? Thanks

Minimally it should say where is it copied from and refer to Pytorch for original file license

Micky774 added 5 commits March 9, 2026 14:06

Updated cmakelist

1dee8f0

Added hsa embed script from pytorch

219cad0

Updated AITER subcommit to feature branch for compatability

484a9e4

Formatting

e9b1484

Updated aiter commit

262b7f2

Micky774 requested review from ipanfilo, wangye805 and wenchenvincent as code owners March 9, 2026 19:21

Micky774 changed the title ~~Zain/aiter hsa embed~~ Utilize AITER HSA embeddings Mar 9, 2026

wangye805 requested changes Mar 10, 2026

View reviewed changes

ipanfilo requested changes Mar 10, 2026

View reviewed changes

transformer_engine/common/ck_fused_attn/generate_aiter_embedded_hsa.py Outdated Show resolved Hide resolved

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

Micky774 added 5 commits March 10, 2026 12:36

Merge branch 'zain/aiter-static' into zain/aiter-hsa-embed

7407678

Updated glob pattern for gfx942 HSA layout

5acdf4e

Simplified var name and added ret check for aiter_build.sh

287ddc3

Streamlined cmake var usage

0f91165

Streamlined HSA handling

38eb092

Micky774 force-pushed the zain/aiter-hsa-embed branch from 1c08539 to 38eb092 Compare March 10, 2026 19:03

Micky774 added 3 commits March 10, 2026 14:09

Merge branch 'zain/aiter-static' into zain/aiter-hsa-embed

7761d31

Merge branch 'zain/aiter-static' into zain/aiter-hsa-embed

58bc670

Simplified cmake vars

249c7f9

ipanfilo reviewed Mar 10, 2026

View reviewed changes

wangye805 requested changes Mar 10, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Utilize AITER HSA embeddings#477

Utilize AITER HSA embeddings#477
Micky774 wants to merge 13 commits intozain/aiter-staticfrom
zain/aiter-hsa-embed

Micky774 commented Mar 9, 2026 •

edited

Loading

Uh oh!

wangye805 left a comment

Uh oh!

Uh oh!

wangye805 Mar 10, 2026

Uh oh!

Micky774 Mar 10, 2026

Uh oh!

Uh oh!

Uh oh!

Micky774 commented Mar 10, 2026

Uh oh!

ipanfilo Mar 10, 2026

Uh oh!

ipanfilo Mar 10, 2026

Uh oh!

wangye805 commented Mar 10, 2026

Uh oh!

wangye805 Mar 10, 2026

Uh oh!

ipanfilo Mar 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Micky774 commented Mar 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!

wangye805 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

wangye805 Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

Micky774 Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Micky774 commented Mar 10, 2026

Uh oh!

ipanfilo Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

ipanfilo Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

wangye805 commented Mar 10, 2026

Uh oh!

wangye805 Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

ipanfilo Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Micky774 commented Mar 9, 2026 •

edited

Loading