Skip to content

[Aiter] upgrade aiter to be compatible with latest TE#209

Closed
GeneDer wants to merge 3 commits into
mainfrom
genesu/upgrade-aiter-2026-01
Closed

[Aiter] upgrade aiter to be compatible with latest TE#209
GeneDer wants to merge 3 commits into
mainfrom
genesu/upgrade-aiter-2026-01

Conversation

@GeneDer

@GeneDer GeneDer commented Jan 15, 2026

Copy link
Copy Markdown
Member

TE recently updated the aiter commit they are using ROCm/TransformerEngine#354 (aiter @ ed4bdea). There is an api change around calling _flash_attn_forward. This PR tries to use the new api offered in aiter.

Signed-off-by: Gene Der Su <e870252314@gmail.com>
Signed-off-by: Gene Der Su <e870252314@gmail.com>
@GeneDer GeneDer marked this pull request as ready for review January 16, 2026 23:06
Copilot AI review requested due to automatic review settings January 16, 2026 23:06

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR updates the aiter dependency to be compatible with the latest TransformerEngine by upgrading to a newer commit that includes API changes to the _flash_attn_forward function.

Changes:

  • Updated aiter package reference to use amd-aiter and a newer commit hash
  • Added new parameters (sink_size, q_descale, k_descale, v_descale) to the _flash_attn_forward call to align with the updated API

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File Description
setup.py Updates aiter dependency to use new package name and commit hash
primus_turbo/pytorch/kernels/attention/attention_csrc_impl.py Adds new parameters to the forward function signature and passes them to _flash_attn_forward

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread setup.py
@GeneDer

GeneDer commented Jan 16, 2026

Copy link
Copy Markdown
Member Author

@wenxie-amd @xiaobochen-amd can you help to merge this? We need it to be able to use the latest TE 🙏

Signed-off-by: Gene Der Su <e870252314@gmail.com>
@xiaobochen-amd

Copy link
Copy Markdown
Collaborator

@GeneDer FYI: this PR also needs to update the AITER version.

#208

@GeneDer

GeneDer commented Jan 19, 2026

Copy link
Copy Markdown
Member Author

duplicated with #208

@GeneDer GeneDer closed this Jan 19, 2026
@GeneDer GeneDer deleted the genesu/upgrade-aiter-2026-01 branch January 19, 2026 15:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants