[Torch] Support mask mod arguments in flex_attention HOP#4570
Open
keshavvinayak01 wants to merge 2 commits into
Open
[Torch] Support mask mod arguments in flex_attention HOP#4570keshavvinayak01 wants to merge 2 commits into
keshavvinayak01 wants to merge 2 commits into
Conversation
keshavvinayak01
force-pushed
the
flex-mask-mod-args
branch
from
2520a91 to
6bec667
Compare
keshavvinayak01
requested review from
Groverkss,
IanWood1,
MaheshRavishankar,
rsuderman,
sjain-stanford and
sommerlukas
keshavvinayak01
force-pushed
the
flex-mask-mod-args
branch
2 times, most recently
from
5ff3fba to
819aea0
Compare
rkayaith
reviewed
May 14, 2026
zjgarvey
requested changes
May 14, 2026
zjgarvey
left a comment
zjgarvey
left a comment
Collaborator
There was a problem hiding this comment.
I'm mostly concerned about the silently dropped arg.
Another thing I'd like to ask is if you would be willing to add verifier logic for at least checking operand count vs. mask_mod_fn arity (or at least something meaningful to cover the new case).
Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
keshavvinayak01
force-pushed
the
flex-mask-mod-args
branch
from
819aea0 to
e780e23
Compare
Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
keshavvinayak01
commented
May 18, 2026
keshavvinayak01
left a comment
keshavvinayak01
left a comment
Contributor
Author
There was a problem hiding this comment.
Thanks for the reviews, please check again.
keshavvinayak01
requested review from
rkayaith and
zjgarvey
and removed request for
rkayaith
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
3 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.
In traced/exported FX graphs, Dynamo represents mask_mod closure captures for the flex_attention HOP as the final mask_mod_other_buffers argument. torch.hop_flex_attention only modelled the mask function symbol, so those captures could not be represented in Torch IR.
This adds trailing variadic
mask_mod_other_buffersoperands totorch.hop_flex_attentionand prints them with explicitmask_mod_other_buffers(...)syntax. The FX importer forwards only exportedmask_mod_other_buffersvalues into those operands.score_mod_other_buffersare named and rejected with an explicit unsupported error instead of being silently dropped.