Skip to content

[C][PyTorch]Make pytorch extensions pure cpp #1754

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
Merged
merged 24 commits into from
May 11, 2025
Merged

[C][PyTorch]Make pytorch extensions pure cpp #1754

merged 24 commits into from
May 11, 2025

Conversation

ksivaman
Copy link
Member

@ksivaman ksivaman commented May 7, 2025 •
edited
Loading

Description

This is a last part in a series of PRs making the framework extensions purely C++.

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 refactoring

Changes

  • Move the remaining cuda attention functionality from PyTorch extensions to core lib, introducing the C APIs as necessary.
  • Fix a bug in the numerics tests.
  • Fix a bug in the fused attention tests.
  • Convert PyTorch extension from cuda to CPP. This is improves compilation speed.
  • mha_fill_kernel has been removed and replaced with nvte_memset.

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

ksivaman added 6 commits May 7, 2025 00:30
@ksivaman
First pass refactor
e9c868e 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Merge branch 'NVIDIA:main' into make_pytorch_extensions_pure_cpp
86e15cf
@ksivaman
first pass
87495f1 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
core compiles
ff32002 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Include cuda dirs
7e4feb4 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Compiles
f368408 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman ksivaman added the 2.4.0 label May 7, 2025
@ksivaman ksivaman marked this pull request as draft May 7, 2025 07:07
@ksivaman ksivaman marked this pull request as ready for review May 7, 2025 17:59
@ksivaman ksivaman requested a review from cyanguwa May 7, 2025 17:59
ksivaman added 7 commits May 7, 2025 20:31
@ksivaman
Fix
4524fc7 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Fix test
b3ca75c 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Merge branch 'NVIDIA:main' into make_pytorch_extensions_pure_cpp
7e5b392
@ksivaman
Move grad outside autocast
69b6998 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Merge branch 'main' into make_pytorch_extensions_pure_cpp
ca24b14
@ksivaman
Fix kv cache
2e69054 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Merge branch 'main' into make_pytorch_extensions_pure_cpp
5ca2d0e
@ksivaman
Copy link
Member Author

ksivaman commented May 9, 2025

/te-ci L0 L1

@ksivaman
Address review comments
5397033 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman ksivaman force-pushed the make_pytorch_extensions_pure_cpp branch from d103923 to 5397033 Compare May 9, 2025 23:12
pre-commit-ci bot and others added 2 commits May 9, 2025 23:13
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
5113e86 
for more information, see https://pre-commit.ci
@ksivaman
Change src file name in cmake
0aa469c 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
move the kernels too
02ccb75 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ksivaman added 5 commits May 9, 2025 23:38
@ksivaman
fix
5b2a7a0 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Move comment
833825f 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Move comments around
f3fe837 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
more movement
709af47 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
move
c31afad 
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
cyanguwa
cyanguwa approved these changes May 10, 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. Thanks for responding so quickly.

@ksivaman
Copy link
Member Author

/te-ci L0 L1

@ksivaman ksivaman merged commit 51cd441 into NVIDIA:main May 11, 2025
29 of 39 checks passed
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
2.4.0
Projects
None yet
Milestone
No milestone
Development

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