Repository navigation
convert _native_multi_head_attention , the fast path of nn.MultiheadAttention - #2899
Open
HuzaifaChaudary wants to merge 2 commits into
Open
HuzaifaChaudary wants to merge 2 commits into
HuzaifaChaudary wants to merge 2 commits into
Conversation
…ttention an eval mode MultiheadAttention traced with nothing needing grad goes through this op , and it had no converter , so any model built on it failed to convert . the mask is read as boolean the way the kernel reads it , so a non zero entry masks the position out .
…ting every case checks the traced graph really holds aten::_native_multi_head_attention , since torch skips the fast path when a parameter needs grad . with the converter taken out all 39 fail on the missing op .
This branch has not been deployed
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
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.
fixes #2311 .
an eval mode
nn.MultiheadAttentionwithbatch_firstand nothing needing grad runs throughaten::_native_multi_head_attention, and there was no converter for it , so any model built on it failed with PyTorch convert function for op '_native_multi_head_attention' not implemented . the only way round it was turning the fast path off in torch .this adds one . it splits the packed in projection into q , k and v , runs them through the same fused or decomposed sdpa the
scaled_dot_product_attentionconverter uses , and returns the weights too when they are asked for , averaged or per head . the kernel reads its mask as boolean , so any non zero entry is masked out , and the converter does the same .the tests cover no mask , key padding , attn mask and both , each weights setting , ios18 and older targets , and a flexible batch and sequence . each one checks the traced graph really has the op , because torch skips the fast path when a parameter needs grad . with the converter taken out all 39 fail on the missing op , with it they pass . the sdpa and transformer tests still pass too , 240 in all , on macos arm64 with torch 2.8.0 .