Enable Runtime Selection of Attention Functions #9639
+316
−179
Conversation
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
…override' entry in transformer_options
…down all the code paths where transformer_options would need to be added
…n to test out attention override later
…d potentially be passed through
…o load SageAttention and FlashAttention if not enabled so that they can be marked as available or not, create registry for available attention
…ve a dropdown with available attention (this is a test node only)
… into attention function
…e to call that, just in case
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.
This PR adds support for choosing the attention to be used for sampling via
optimized_attention_overrideintransformer_optionsin all natively-supported models. This opens up a lot of doors for attention tricks/scheduling/different attention in different blocks without massive hacks/monkey patches.All available attention types are also registered so it can be tracked/chosen easily in the future - flash attn/sage attention not being available will only cause immediate exit if the corresponding
--use-flash-attentionor--use-sage-attentionarguments are used. Otherwise, the attention types will still be registered but only used upon request in code.More changes are coming in the future regarding models tracking the current block index in
transformer_options.Changes went through two rounds of QA + my personal testing. No observable slowdowns were noticed.