Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Ipanfilo/fa support #96

Open
wants to merge 11 commits into
base: dev
Choose a base branch
from
Open

Ipanfilo/fa support #96

wants to merge 11 commits into from

Conversation

ipanfilo
Copy link
Contributor

@ipanfilo ipanfilo commented Dec 4, 2024

Description

This is work-in-progress review for ROCm FlashAttn2 support

Fixes #9693

Type of change

  • New feature (non-breaking change which adds functionality)

Changes

  • Enable FA specific code on ROCm platforms
  • Cherry-pick upstream TE changes needed for proper FA working
  • Modify test scripts to accommodate new FA

BestJuly and others added 5 commits November 28, 2024 18:28
fix an argument issue when flash_attn>=2.5.7

Signed-off-by: Li Tao <lit@nvidia.com>
Co-authored-by: Li Tao <lit@nvidia.com>
Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>
Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ipanfilo ipanfilo mentioned this pull request Dec 5, 2024
13 tasks
@@ -12,7 +12,7 @@
from test_fused_attn import ModelConfig
from transformer_engine.pytorch.attention import (
_flash_attn_2_plus,
_flash_attn_2_3_plus,
_flash_attn_2_6_plus,
Copy link
Collaborator

Choose a reason for hiding this comment

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

It seems that this is not used in this file?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Oh, yes. It left from pre IFU 1.11 commit where THD was only skipped with 2.6+

Copy link
Collaborator

Choose a reason for hiding this comment

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

Oh, also, I think we need to merge the updates from the dev branch (now with IFU 1.11) to make sure there is no breakage from that.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

It has those changes

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.

5 participants