Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Hunt correctness for extend attention + RPE #504

Draft
wants to merge 10 commits into
base: main
Choose a base branch
from
Draft

Hunt correctness for extend attention + RPE #504

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
e6ef6d7
[Wave] Refactor style and fix test for RPE
raikonenfnu Feb 13, 2025
d1168c9
[Wave] Add RPE to Extend Attention + Fix shapes for vanilla RPE
raikonenfnu Feb 13, 2025
8726a56
Hunt correctness for extend attention + RPE
nicolasvasilache Feb 14, 2025
8fae0d4
Thread max_rpe_context_lengths were appropriate
ftynse Feb 14, 2025
d6d5577
thread more
ftynse Feb 14, 2025
89d01e8
debug rpe
nicolasvasilache Feb 14, 2025
4757d8f
Add debug_rpe to dump out the RPE mask from the kernel
nicolasvasilache Feb 14, 2025
259edc0
Minimize and dump IR
nicolasvasilache Feb 14, 2025
6bf801e
debug step
nicolasvasilache Feb 14, 2025
22192e5
Allow running without pytest
nicolasvasilache Feb 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions iree/turbine/kernel/ops/wave_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1233,7 +1233,9 @@ def align_index(self, constraints: list["Constraint"]) -> None:
from ..wave.utils import align_index_vars, is_shared_mem_access

if is_shared_mem_access(self):
print(f"MMA before {self.index} with constraints {constraints}")
self.index = align_index_vars(self.index, constraints)
print(f"MMA after {self.index}")

def transform_index_backwards(
self, index: dict[IndexSymbol, IndexSequence], arg: fx.Node
Expand Down
Loading