Skip to content

[PyTorch] Add sliding window support to FlashAttention - #551

Merged
ptrendx merged 21 commits into
NVIDIA:mainfrom
cyanguwa:fa/sliding_window
Dec 16, 2023
Merged

[PyTorch] Add sliding window support to FlashAttention#551
ptrendx merged 21 commits into
NVIDIA:mainfrom
cyanguwa:fa/sliding_window

Conversation

@cyanguwa

@cyanguwa cyanguwa commented Dec 6, 2023

Copy link
Copy Markdown
Collaborator

This PR only makes changes on the PyTorch side. It

  • integrates flash-attn 2.3+ sliding window attention to TransformerLayer, MultiHeadAttention, DotProductAttention and FlashAttention
  • adds unit tests to compare against UnfusedDotProductAttention arbitrary mask, generated based on the window size
  • adds a use_unfused_attention flag and exception when none of the three DPA backends are available
  • adds more determinism control in the backend selection, in particular, the filter that fused attention arbitrary backend is non-deterministic on non-sm90 architectures because it doesn't have a workspace optimization path

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

cyanguwa commented Dec 6, 2023

Copy link
Copy Markdown
Collaborator Author

/te-ci

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

cyanguwa commented Dec 6, 2023

Copy link
Copy Markdown
Collaborator Author

/te-ci

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

cyanguwa commented Dec 6, 2023

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa cyanguwa closed this Dec 7, 2023
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa cyanguwa reopened this Dec 8, 2023
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

cyanguwa commented Dec 8, 2023

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch

1 similar comment
@cyanguwa

cyanguwa commented Dec 8, 2023

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch

@cyanguwa

cyanguwa commented Dec 8, 2023

Copy link
Copy Markdown
Collaborator Author

Pipeline 11331817

@cyanguwa

cyanguwa commented Dec 11, 2023

Copy link
Copy Markdown
Collaborator Author

With newer cuDNN, pipeline 11412889 is green!

Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch

Comment thread tests/pytorch/test_fused_attn.py
Comment thread tests/pytorch/test_fused_attn.py
Comment thread tests/pytorch/test_sanity.py Outdated
Comment thread transformer_engine/pytorch/attention.py Outdated
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Comment thread transformer_engine/pytorch/attention.py Outdated
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch

@ptrendx ptrendx left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

LGTM, thanks!

@ashvinnihalani

Copy link
Copy Markdown

Also want to add a comment that what happens when we want to use grouped query attention with unfused attention. Right now it seems like it errors out.

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