[DO NOT MERGE] cudagraph: use cuda dynamic parallelism to dispatch kernels #288
Conversation
This file contains 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
|
This is really fascinating! Curious to know about the alternate approach. Let me know if I can help in anyway. Thanks! |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Background
FlashInfer will dispatch to different kernels depending on input shapes, which is not cuda-graph friendly at the moment.
CUDA 12.4 announced conditional nodes in CUDA Graphs, and it will take some time for PyTorch to fully support this feature (capture conditional nodes).
Dynamic Parallelism
In #187 , @AgrawalAmey propose to use CUDA's Dynamic Parallelism feature to dispatch to different child kernels from a parent launcher kernel. We don't need to capture conditional nodes in this way.
Discussion
This PR implements the dynamic parallelism launcher, however, we need to enable
-rdc(relocatable-device-code) in compiler options to support dynamic parallelism, which introduce some non-negligible runtime overhead for small kernels (there are some discussions in nvidia forum).Considering the potential negative consequence of enabling rdc (e.g. performance degradation), I'll hold this PR and turn to use alternative approach (which do not require dispatching) to support CUDA Graph. We won't merge this PR at the moment and we welcome discussions.