Skip to content

Added generated test files for FlexAttention#149

Draft
keshavvinayak01 wants to merge 5 commits into
iree-org:mainfrom
keshavvinayak01:users/keshavvinayak01/flex-attention-torchops-tests
Draft

Added generated test files for FlexAttention#149
keshavvinayak01 wants to merge 5 commits into
iree-org:mainfrom
keshavvinayak01:users/keshavvinayak01/flex-attention-torchops-tests

Conversation

@keshavvinayak01
Copy link
Copy Markdown
Contributor

@keshavvinayak01 keshavvinayak01 commented Jan 16, 2026

Adding torch_ops test for FlexAttention, as requested in iree-org/iree#23292.
(Not to be merged until the above goes in)

@keshavvinayak01 keshavvinayak01 marked this pull request as draft January 16, 2026 10:17
Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
keshavvinayak01 added a commit to iree-org/iree that referenced this pull request Apr 22, 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>
benvanik pushed a commit to iree-org/iree 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 to iree-org/iree 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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant