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.

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

clean CP implementation for flash attention and cuDNN 9.6 #1387

Merged
merged 24 commits into from
Jan 8, 2025
Merged

clean CP implementation for flash attention and cuDNN 9.6 #1387

merged 24 commits into from
Jan 8, 2025
+442 −436

Conversation

xrennvidia
Copy link
Collaborator

@xrennvidia xrennvidia commented Dec 30, 2024 •
edited by cyanguwa
Loading

Description

  • Flash attention: only use varlen kernels for THD format
  • cuDNN >=9.6: fix implementation for packed format of softmax_lse, enabled THD+GQA, enabled SWA with KV all_gather communications.

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

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

xrennvidia and others added 21 commits November 28, 2024 19:50
@xrennvidia
make pad_between_seqs check do not consider padding at the end
057de06 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
change CP THD test to make it consider 0-length sequence
4bcc1ce 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
minor change to flash func name
62af46f 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
only use varlen func of flash attention while qkv_format is THD
fdc83fe 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
try to converge code of flash and fused attentions
a7e14bf 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
Merge branch 'main' into xren/cp_optim
cc2f3bc
@xrennvidia
Merge branch 'main' into xren/cp_optim
6041e8d
@xrennvidia
Merge branch 'main' into xren/cp_optim
685b08c
@xrennvidia
Merge branch 'main' into xren/cp_optim
d69198a
@xrennvidia
Merge branch 'main' into xren/cp_optim
7157ad7
@xrennvidia
Merge branch 'main' into xren/cp_optim
2bccdd0
@xrennvidia
fix bwd compute with P2P
0e62ee1 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
remove redundant out_per_step view
1706ec4 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
enable cudnn>9.6 and THD+GQA
620f86c 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
enable CP with FusedAttn+SWA+All_Gather
9ca6f8d 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
enable CP with FusedAttn+SWA+All_Gather
7131139 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
code cleaning for cu_seqlens
da50f5b 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
d9443d4 
for more information, see https://pre-commit.ci
@xrennvidia
fix some pylint error
09b4f4f 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
minor import change for pylint
1b0eac9 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
more fix for pylint
0dc2553 
Signed-off-by: Xiaowei Ren <xren@nvidia.com>
@xrennvidia
Copy link
Collaborator Author

/te-ci pytorch L1

@xrennvidia xrennvidia requested a review from cyanguwa December 31, 2024 08:19
cyanguwa
cyanguwa approved these changes Jan 6, 2025
View reviewed changes
Copy link
Collaborator

@cyanguwa cyanguwa left a comment

Choose a reason for hiding this comment

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

LGTM

transformer_engine/pytorch/attention.py Show resolved Hide resolved
@cyanguwa
Copy link
Collaborator

cyanguwa commented Jan 6, 2025

/te-ci pytorch L0

@xrennvidia
Copy link
Collaborator Author

/te-ci pytorch L1

@xrennvidia xrennvidia merged commit 560bccf into NVIDIA:main Jan 8, 2025
14 checks passed
@xrennvidia xrennvidia deleted the xren/cp_optim branch January 8, 2025 18:13
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@cyanguwa cyanguwa cyanguwa approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@xrennvidia @cyanguwa