Fix flash-attention build env in CI deps Dockerfile by VeeraRajasekhar · Pull Request #565 · ROCm/TransformerEngine

VeeraRajasekhar · 2026-04-27T20:35:29Z

…CI image tag

Description

Please include a brief summary of the changes, relevant motivation and context.

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

VeeraRajasekhar · 2026-04-27T20:47:48Z

CI Run: https://github.com/ROCm/TransformerEngine/actions/runs/25018699393

sudhu2k · 2026-04-28T01:18:05Z

LGTM.

ipanfilo · 2026-04-28T04:07:11Z


-# AITER (submodule: composable_kernel / CK). Submodule update after checkout aligns CK with this commit.
+# Install AITER
 RUN git clone --recursive https://github.com/ROCm/aiter.git \


You don't need --recursive if later calling git submodule update --init --recursive. Doing that you checkout submodules twice.
And if doing checkout later, makes sense to add --no-checkout to clone

ipanfilo · 2026-04-28T04:08:40Z

Remove PIP_FIND_LINKS and PIP_PRE. The used AITER commit does not depend of FlyDSL and there is no FlyDSL distribution for Python 3.11 anyway

ipanfilo · 2026-04-28T04:11:39Z

 {
  "docker_images": {
-    "default": "registry-sc-harbor.amd.com/framework/te-ci:rocm-7.2_ubuntu22.04_py3.11_pytorch_release-2.8_08d38866_jax_0.8.0_fa_2.8.1_aiter_77455e3ecf",
+    "default": "registry-sc-harbor.amd.com/framework/te-ci:rocm-7.2_ubuntu22.04_py3.11_pytorch_release-2.8_08d38866_jax_0.8.0_fa_2.8.1_aiter_77455e3ecf_v1",


For GHA the name can be reused, AFAIK - there is no image caching there

No, I already pushed an image with "default": "registry-sc-harbor.amd.com/framework/te-ci:rocm-7.2_ubuntu22.04_py3.11_pytorch_release-2.8_08d38866_jax_0.8.0_fa_2.8.1_aiter_77455e3ecf", and later created a new image.

Can that image not be overwritten?

It should be possible to overwritten. In the works case old one can be deleted. With Jenkins we had to use uniq tags because there was image caching on workers that did not look at hash but only name/tag

Updated to use the older tag

Micky774 · 2026-04-28T17:35:33Z

+    && export GPU_ARCHS="gfx950;gfx942" \
+    && FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE FLASH_ATTENTION_SKIP_CK_BUILD=FALSE python setup.py install \


Why an export instead of feeding it as an env var before the python setup.py install?

It is probably not needed at all for installation - nothing is compiled there

Micky774

@VeeraRajasekhar @sudhu2k can either of you confirm that the image as it is now passes the failing tests for MXFP4 we're seeing on dev?

Edit: The current running CI for this PR is using the updated image source.

sudhu2k · 2026-04-29T05:26:31Z

@VeeraRajasekhar @sudhu2k can either of you confirm that the image as it is now passes the failing tests for MXFP4 we're seeing on dev?

Edit: The current running CI for this PR is using the updated image source.

https://github.com/ROCm/TransformerEngine/actions/runs/25014847020/job/73259852413
The level 3 tests all passed along with the mxfp4 GEMM tests when I ran it yesterday before merging the MXFP4 PR with the updated image.

…mage tag

VeeraRajasekhar self-assigned this Apr 27, 2026

VeeraRajasekhar requested review from ipanfilo, wangye805 and wenchenvincent as code owners April 27, 2026 20:35

VeeraRajasekhar force-pushed the veergopu/update_ci_image branch from 8989d90 to 2f20c17 Compare April 27, 2026 20:39

VeeraRajasekhar added the ci-level 1 CI test level 1 label Apr 27, 2026

ipanfilo requested changes Apr 28, 2026

View reviewed changes

VeeraRajasekhar requested a review from ipanfilo April 28, 2026 16:59

Micky774 reviewed Apr 28, 2026

View reviewed changes

wangye805 requested changes Apr 28, 2026

View reviewed changes

Comment thread .github/scripts/Dockerfile.ci.deps Outdated

wangye805 requested changes Apr 28, 2026

View reviewed changes

Comment thread .github/scripts/Dockerfile.ci.deps Outdated

VeeraRajasekhar requested review from Micky774 and wangye805 April 28, 2026 20:23

Micky774 approved these changes Apr 28, 2026

View reviewed changes

wangye805 approved these changes Apr 28, 2026

View reviewed changes

VeeraRajasekhar changed the title ~~Fix flash-attention build env in CI deps Dockerfile; bump default TE image tag~~ Fix flash-attention build env in CI deps Dockerfile Apr 29, 2026

ipanfilo approved these changes Apr 29, 2026

View reviewed changes

VeeraRajasekhar added 3 commits April 29, 2026 11:50

Fix flash-attention build in CI deps Dockerfile; bump default TE CI i…

a4c0942

…mage tag

Addressed reviews

a30316e

Addressed reviews

ba241bf

VeeraRajasekhar force-pushed the veergopu/update_ci_image branch from 892dc54 to ba241bf Compare April 29, 2026 16:51

VeeraRajasekhar merged commit 732e5ab into dev Apr 29, 2026
3 checks passed

		&& export GPU_ARCHS="gfx950;gfx942" \
		&& FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE FLASH_ATTENTION_SKIP_CK_BUILD=FALSE python setup.py install \

Conversation

VeeraRajasekhar commented Apr 27, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

VeeraRajasekhar commented Apr 27, 2026

Uh oh!

sudhu2k commented Apr 28, 2026

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!

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 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sudhu2k commented Apr 29, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Micky774 left a comment •

edited

Loading