File tree Expand file tree Collapse file tree 1 file changed +10
-2
lines changed
extension/llm/modules/test Expand file tree Collapse file tree 1 file changed +10
-2
lines changed Original file line number Diff line number Diff line change 10
10
11
11
import torch
12
12
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
13
15
14
16
from executorch .extension .export_util .utils import save_pte_program
15
17
from executorch .extension .llm .modules .kv_cache import KVCache as InferenceKVCache
@@ -171,7 +173,10 @@ def test_kv_cache_executorch(self):
171
173
_check_ir_validity = False ,
172
174
),
173
175
)
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 )
175
180
176
181
runtime = Runtime .get ()
177
182
program = runtime .load_program (et_program .buffer )
@@ -192,7 +197,10 @@ def test_kv_cache_executorch_from_file(self):
192
197
_check_ir_validity = False ,
193
198
),
194
199
)
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 )
196
204
197
205
with tempfile .TemporaryDirectory () as tempdir :
198
206
pte_path = save_pte_program (et_program , "test_et_kv_cache" , tempdir )
You can’t perform that action at this time.
0 commit comments