Skip to content

Commit 3b475e3

Browse files
authored
[llama-mm] Fix AOTI test for attention (#6915)
Summary: Disable `reorder_for_peak_memory` because it moves `_local_dense_scalar` codegen to after subgraphs. Test Plan: As titled. ``` RUN_SKIPPED=1 pytorch -m unittest extension.llm.modules.test.test_attention -k test_attention_aoti ``` Need to address the flaky test later. Reviewers: Subscribers: Tasks: Tags:
1 parent 4f9ae32 commit 3b475e3

File tree

1 file changed

+7
-2
lines changed

1 file changed

+7
-2
lines changed

extension/llm/modules/test/test_attention.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -156,7 +156,9 @@ def test_attention_export(self):
156156

157157
assert_close(et_res, tt_res)
158158

159-
@unittest.skip(reason="TODO(T207740932): test is flaky")
159+
@unittest.skipIf(
160+
int(os.getenv("RUN_SKIPPED", 0)) < 1, reason="TODO(T207740932): test is flaky"
161+
)
160162
def test_attention_aoti(self):
161163
# Self attention.
162164

@@ -168,7 +170,10 @@ def test_attention_aoti(self):
168170
self.et_mha,
169171
args=(self.x, self.x),
170172
kwargs={"input_pos": self.input_pos},
171-
options={"aot_inductor.package": True},
173+
options={
174+
"aot_inductor.package": True,
175+
"reorder_for_peak_memory": False,
176+
},
172177
dynamic_shapes=self.dynamic_shapes,
173178
)
174179
with tempfile.TemporaryDirectory() as tempdir:

0 commit comments

Comments
 (0)