Skip to content

[common/PyTorch] Add FusedAttention support for SWA (left, right) #1369

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.

Sign up for GitHub

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

Jump to bottom
Closed
cyanguwa wants to merge 18 commits into from
Closed

[common/PyTorch] Add FusedAttention support for SWA (left, right) #1369

cyanguwa wants to merge 18 commits into from

Conversation

@cyanguwa
Copy link
Collaborator

@cyanguwa cyanguwa commented Dec 12, 2024 •
edited
Loading

Description

This PR adds SWA (left, right) support to the FusedAttention backend for {bshd, sbhd, thd} formats and {no_mask, causal, causal_bottom_right, padding, padding_causal, padding_causal_bottom_right} masks.

Previously, the support was only for SWA (left, 0), {bshd, sbhd} formats, and {causal} mask.

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 refractor

Changes

Please list the changes introduced in this PR:

  • Adds FusedAttention support for SWA (left, right)

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

cyanguwa and others added 6 commits December 11, 2024 21:30
@cyanguwa
WIP: add support for SWA (left,0) + THD/BSHD/SBHD + padding + CM/BRCM
165f99c 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
e2d9ffe 
for more information, see https://pre-commit.ci
@cyanguwa
enable more support
8572c1f 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
681ffbe 
for more information, see https://pre-commit.ci
@cyanguwa
WIP: fix up swa
956570f 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
8d17e10 
for more information, see https://pre-commit.ci
@cyanguwa cyanguwa marked this pull request as ready for review December 15, 2024 02:07
@cyanguwa
Copy link
Collaborator Author

cyanguwa commented Dec 15, 2024

/te-ci pytorch L0

cyanguwa and others added 4 commits December 16, 2024 01:25
@cyanguwa
add left_bound/right_bound
9a09edb 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
277dd60 
for more information, see https://pre-commit.ci
@cyanguwa
tweak tests
fc6e338 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
dff7e09 
for more information, see https://pre-commit.ci
@cyanguwa cyanguwa changed the title [common/PyTorch] Add support for SWA (left,0) + padding + CM/BRCM [common/PyTorch] Add support for SWA (left,right) Dec 16, 2024
@cyanguwa cyanguwa changed the title [common/PyTorch] Add support for SWA (left,right) [common/PyTorch] Add support for SWA (left, right) Dec 16, 2024
@cyanguwa cyanguwa changed the title [common/PyTorch] Add support for SWA (left, right) [common/PyTorch] Add FusedAttention support for SWA (left, right) Dec 16, 2024
cyanguwa and others added 7 commits December 16, 2024 05:03
@cyanguwa
fix C swa and tests
e64a291 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
fix get_swa_mask
a8ca89c 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
dd9159b 
for more information, see https://pre-commit.ci
@cyanguwa
fix lint
8c4d836 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
6f677da 
for more information, see https://pre-commit.ci
@cyanguwa
Merge branch 'main' into swa_padding_brcm
6a8e073
@cyanguwa
Merge branch 'main' into swa_padding_brcm
0998f9e
@cyanguwa cyanguwa marked this pull request as draft December 17, 2024 10:19
@cyanguwa
Merge branch 'main' into swa_padding_brcm
cd529f9 
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa cyanguwa added the 2.3.0 label Mar 18, 2025
@cyanguwa cyanguwa removed the 2.3.0 label Apr 10, 2025
sudhakarsingh27 added a commit to sudhakarsingh27/TransformerEngine that referenced this pull request Dec 4, 2025
@sudhakarsingh27
SWA (left, right) with FusedAttention changes cherry-picked from NVID…
33e2e47 
…IA#1369

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
@sudhakarsingh27 sudhakarsingh27 mentioned this pull request Dec 4, 2025
Open
28 tasks
@cyanguwa
Copy link
Collaborator Author

cyanguwa commented Dec 4, 2025

Superseded by #2477.

@cyanguwa cyanguwa closed this Dec 4, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@sudhakarsingh27 sudhakarsingh27 Awaiting requested review from sudhakarsingh27

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@cyanguwa