Skip to content

Commit e297c9b

Browse files
committed
Fix test
1 parent 4ee95d3 commit e297c9b

File tree

1 file changed

+10
-2
lines changed

1 file changed

+10
-2
lines changed

extension/llm/modules/test/test_kv_cache.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010

1111
import torch
1212
from executorch.exir import EdgeCompileConfig, to_edge
13+
from executorch.exir.capture._config import ExecutorchBackendConfig
14+
from executorch.exir.passes.init_mutable_pass import InitializedMutableBufferPass
1315

1416
from executorch.extension.export_util.utils import save_pte_program
1517
from executorch.extension.llm.modules.kv_cache import KVCache as InferenceKVCache
@@ -171,7 +173,10 @@ def test_kv_cache_executorch(self):
171173
_check_ir_validity=False,
172174
),
173175
)
174-
et_program = edge_program.to_executorch()
176+
et_config = ExecutorchBackendConfig(
177+
passes=[InitializedMutableBufferPass(["cache_pos"])],
178+
)
179+
et_program = edge_program.to_executorch(config=et_config)
175180

176181
runtime = Runtime.get()
177182
program = runtime.load_program(et_program.buffer)
@@ -192,7 +197,10 @@ def test_kv_cache_executorch_from_file(self):
192197
_check_ir_validity=False,
193198
),
194199
)
195-
et_program = edge_program.to_executorch()
200+
et_config = ExecutorchBackendConfig(
201+
passes=[InitializedMutableBufferPass(["cache_pos"])],
202+
)
203+
et_program = edge_program.to_executorch(config=et_config)
196204

197205
with tempfile.TemporaryDirectory() as tempdir:
198206
pte_path = save_pte_program(et_program, "test_et_kv_cache", tempdir)

0 commit comments

Comments
 (0)