Add Triton support for KNN and Nearest; remove deprecated torch.ji.script#264
Open
denix56 wants to merge 48 commits intorusty1s:masterfrom
Open
Add Triton support for KNN and Nearest; remove deprecated torch.ji.script#264denix56 wants to merge 48 commits intorusty1s:masterfrom
denix56 wants to merge 48 commits intorusty1s:masterfrom
Conversation
Use sorted batch pointers for Triton masks
Fix Triton pairwise grid sizing
Add Triton-based variants of KNN and Nearest search algorithms
Compute split-N tile counts via heuristics
* Compute split-N tile counts via heuristics * Add Triton support * Add Triton support
* Compute split-N tile counts via heuristics * Add Triton support * Add Triton support * Add Triton support
remove triton from dependencies
* Compute split-N tile counts via heuristics * Add Triton support * Add Triton support * Add Triton support * Fix flake8 errors remove triton from dependencies
remove triton from dependencies
* Compute split-N tile counts via heuristics * Add Triton support * Add Triton support * Add Triton support * Fix flake8 errors remove triton from dependencies * Remove meaningless comments
Fix flake8 errors
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #264 +/- ##
===========================================
- Coverage 97.32% 37.64% -59.68%
===========================================
Files 11 17 +6
Lines 224 704 +480
===========================================
+ Hits 218 265 +47
- Misses 6 439 +433 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Gate torch.compile in tests for windows
Fix flake8 errors
Author
|
Some benchmarks |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Hi everyone,
I have implemented Triton versions of KNN and Nearest that achieve up to 2x times faster performance (especially on bigger data). The triton implementation is behind a flag
use_triton.This implementation will be particularly useful for NNs that usually have fixed shapes, providing a noticeable speedup.
Here are the metrics (respective benchmarks are included in project):
benchmarks.rtf
Also I have removed now deprecated torch.jit.script and replaced operator registration with TORCH_LIBRARY. Moreover, I have added fake kernels for every op in the library, allowing to use them in with torch.compile.