[LinalgExt] Rewriter for Torch::HigherOrderFlexAttentionOp -> LinalgExt::OnlineAttentionOp#23292
Merged
keshavvinayak01 merged 7 commits intoApr 22, 2026
Conversation
This was referenced Jan 27, 2026
Collaborator
|
@keshavvinayak01 moving PRs around makes it hard to track what is new and what has been up for a while. It is disruptive for reviewers. Can we keep this a bit more stable? |
Contributor
Author
|
@Groverkss Could you close reviews on this? |
Convert torch.hop_flex_attention -> iree_linalg_ext.online_attention with inlined score/mask modification functions. The mask_mod and score_mod function bodies are inlined directly into the score modification region (no func.call, no separate mask tensor), enabling fusion during attention decomposition and proper tiling. Also fixes: - IndexOp verifier to accept OnlineAttentionOp as parent - OnlineAttentionOp::build scale/mask parameter swap - applyPostQKMatmulElementwise to convert iree_linalg_ext.index -> linalg.index when cloning the score region during decomposition Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
keshavvinayak01
force-pushed
the
personal/users/keshavvinayak01/linalgext-torch-rewrite-flexattention
branch
from
507e0ef to
8efbb1d
Compare
Compute scale = rsqrt(head_dim) at runtime via tensor.dim + math.rsqrt when the scale is not a constant float, instead of requiring a static head dimension. Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
keshavvinayak01
changed the title
keshavvinayak01
requested review from
IanWood1,
MaheshRavishankar,
Max191 and
hanhanW
as code owners
keshavvinayak01
commented
Apr 10, 2026
Contributor
Author
keshavvinayak01
left a comment
keshavvinayak01
left a comment
There was a problem hiding this comment.
Re-opening with the rewriter converting to online_attention directly instead of AttentionOp. We might also use this torch op in fusilli cc @sjain-stanford so pulling you in for reviews.
MaheshRavishankar
approved these changes
Apr 14, 2026
Collaborator
MaheshRavishankar
left a comment
MaheshRavishankar
left a comment
There was a problem hiding this comment.
Overall looks fine to me.
@rsuderman can you review this PR and the follow ups on this if no one gets to it.
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
keshavvinayak01
deleted the
personal/users/keshavvinayak01/linalgext-torch-rewrite-flexattention
branch
benvanik
pushed a commit
that referenced
this pull request
Apr 24, 2026
…xt::OnlineAttentionOp (#23292) Rewriter pattern for torch.hop_flex_attention -> iree_linalg_ext.online_attention I ran the entire flex_attention_hop implementation with randomised input tensors, (Also see llvm/torch-mlir#4366) through aot.export and compared against eager mode, and I noticed no accuracy losses (On CPU) Test: [Torch ops test PR ](iree-org/iree-test-suites#149) --------- Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
jerryyin
pushed a commit
that referenced
this pull request
May 7, 2026
…xt::OnlineAttentionOp (#23292) Rewriter pattern for torch.hop_flex_attention -> iree_linalg_ext.online_attention I ran the entire flex_attention_hop implementation with randomised input tensors, (Also see llvm/torch-mlir#4366) through aot.export and compared against eager mode, and I noticed no accuracy losses (On CPU) Test: [Torch ops test PR ](iree-org/iree-test-suites#149) --------- Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
RattataKing
pushed a commit
to RattataKing/iree
that referenced
this pull request
May 27, 2026
…xt::OnlineAttentionOp (iree-org#23292) Rewriter pattern for torch.hop_flex_attention -> iree_linalg_ext.online_attention I ran the entire flex_attention_hop implementation with randomised input tensors, (Also see llvm/torch-mlir#4366) through aot.export and compared against eager mode, and I noticed no accuracy losses (On CPU) Test: [Torch ops test PR ](iree-org/iree-test-suites#149) --------- Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
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.
Following the discussion from #22441.
I ran the entire flex_attention_hop implementation with randomised input tensors, (Also see llvm/torch-mlir#4366) through aot.export and compared against eager mode, and I noticed no accuracy losses (On CPU)
Test: Torch ops test PR