[PyTorch] Add sliding window support to FlashAttention by cyanguwa · Pull Request #551 · NVIDIA/TransformerEngine

cyanguwa · 2023-12-06T00:26:45Z

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>

cyanguwa · 2023-12-06T00:27:15Z

/te-ci

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

cyanguwa · 2023-12-06T00:35:12Z

/te-ci

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

cyanguwa · 2023-12-06T20:26:47Z

/te-ci pytorch

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

cyanguwa · 2023-12-08T01:13:48Z

/te-ci pytorch

cyanguwa · 2023-12-08T01:29:38Z

/te-ci pytorch

cyanguwa · 2023-12-08T01:47:52Z

Pipeline 11331817

cyanguwa · 2023-12-11T21:40:49Z

With newer cuDNN, pipeline 11412889 is green!

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

cyanguwa · 2023-12-13T23:46:52Z

/te-ci pytorch

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

cyanguwa · 2023-12-15T20:32:49Z

/te-ci pytorch

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

cyanguwa · 2023-12-15T22:43:15Z

/te-ci pytorch

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

cyanguwa · 2023-12-15T23:14:34Z

/te-ci pytorch

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

cyanguwa · 2023-12-15T23:55:27Z

/te-ci pytorch

ptrendx

LGTM, thanks!

ashvinnihalani · 2024-03-07T04:28:11Z

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.

cyanguwa added 2 commits December 6, 2023 00:07

add sliding window to FA
69f568e
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fix forward logic
8f7121f
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

cyanguwa added 3 commits December 6, 2023 00:59

change bert test to causal as unfused does not support padding
37f1469
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fix FlashAttention for v2-2.3 versions
e278579
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Merge branch 'main' into fa/sliding_window
37c1caf

verify FA swa works
49e147e
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa closed this Dec 7, 2023

cyanguwa force-pushed the fa/sliding_window branch from 49e147e to 32db392Compare December 7, 2023 23:41

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

cyanguwa reopened this Dec 8, 2023

cyanguwa added 2 commits December 8, 2023 00:54

fix mask related restrictions and duplicate code after merge
97f807f
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

fix swa test
603cb52
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa requested review from ksivaman, ptrendx and timmoon10 December 11, 2023 22:36

Merge branch 'main' into fa/sliding_window
ce15516
Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>

akoumpa mentioned this pull request Dec 15, 2023
Sliding window attention/akoumparouli NVIDIA/Megatron-LM#633
Merged

ptrendx reviewed Dec 15, 2023
View reviewed changes

Comment threadtests/pytorch/test_fused_attn.py

ptrendx reviewed Dec 15, 2023
View reviewed changes

Comment threadtests/pytorch/test_fused_attn.py

ptrendx reviewed Dec 15, 2023
View reviewed changes

Comment threadtests/pytorch/test_sanity.py Outdated

ptrendx reviewed Dec 15, 2023
View reviewed changes

Comment threadtransformer_engine/pytorch/attention.py Outdated

cyanguwa added 4 commits December 15, 2023 11:51

add docstring for get_swa func
e2fcf1d
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

move repeated code into a function
0f6b235
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Merge branch 'main' into fa/sliding_window
f91d80d

revert mask change
67dd2d9
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

add determinism filter and fix FA warning message
576112d
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

ptrendx reviewed Dec 15, 2023
View reviewed changes

Comment threadtransformer_engine/pytorch/attention.py Outdated

cyanguwa added 2 commits December 15, 2023 22:37

add message for determinism filter
9c7de54
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Merge branch 'main' into fa/sliding_window
c8afd21

simplify check_set_window_size()
53bf1f5
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa added 2 commits December 15, 2023 15:53

fix check_set_window_size in transformer layers
24252a2
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

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

ptrendx approved these changes Dec 16, 2023
View reviewed changes

ptrendx merged commit 27aa609 into NVIDIA:mainDec 16, 2023

Marks101 mentioned this pull request Jan 25, 2024
[PyTorch] FlashAttention: causal masking enforced in cross attention due to sliding window attention #629
Closed

cyanguwa deleted the fa/sliding_window branch February 21, 2024 23:59

Uh oh!

Conversation

cyanguwa commented Dec 6, 2023• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cyanguwa commented Dec 6, 2023

Uh oh!

cyanguwa commented Dec 6, 2023

Uh oh!

cyanguwa commented Dec 6, 2023

Uh oh!

cyanguwa commented Dec 8, 2023

Uh oh!

cyanguwa commented Dec 8, 2023

Uh oh!

cyanguwa commented Dec 8, 2023

Uh oh!

cyanguwa commented Dec 11, 2023• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cyanguwa commented Dec 13, 2023

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cyanguwa commented Dec 15, 2023

Uh oh!

Uh oh!

cyanguwa commented Dec 15, 2023

Uh oh!

cyanguwa commented Dec 15, 2023

Uh oh!

cyanguwa commented Dec 15, 2023

Uh oh!

ptrendx left a comment

Choose a reason for hiding this comment

Uh oh!

ashvinnihalani commented Mar 7, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

cyanguwa commented Dec 6, 2023•
edited
Loading

cyanguwa commented Dec 11, 2023•
edited
Loading