From ce2a9eba2c2a8d4969e9c24bf687603911ac68bb Mon Sep 17 00:00:00 2001 From: Mengwei Liu Date: Fri, 15 Nov 2024 16:15:42 -0800 Subject: [PATCH] [llama-mm] Fix AOTI test for attention 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: --- extension/llm/modules/test/test_attention.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/extension/llm/modules/test/test_attention.py b/extension/llm/modules/test/test_attention.py index f4e4b8c670c..cda9becd69f 100644 --- a/extension/llm/modules/test/test_attention.py +++ b/extension/llm/modules/test/test_attention.py @@ -156,7 +156,9 @@ def test_attention_export(self): assert_close(et_res, tt_res) - @unittest.skip(reason="TODO(T207740932): test is flaky") + @unittest.skipIf( + int(os.getenv("RUN_SKIPPED", 0)) < 1, reason="TODO(T207740932): test is flaky" + ) def test_attention_aoti(self): # Self attention. @@ -168,7 +170,10 @@ def test_attention_aoti(self): self.et_mha, args=(self.x, self.x), kwargs={"input_pos": self.input_pos}, - options={"aot_inductor.package": True}, + options={ + "aot_inductor.package": True, + "reorder_for_peak_memory": False, + }, dynamic_shapes=self.dynamic_shapes, ) with tempfile.TemporaryDirectory() as tempdir: