Skip to content

Commit d0192e5

Browse files
committed
add llama sdpa to generation script
1 parent 99e1ae1 commit d0192e5

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

examples/models/llama2/runner/generation.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,11 @@
1919
Message,
2020
Tokenizer,
2121
)
22+
23+
from executorch.extension.pybindings import portable_lib # noqa
24+
# Note: import this after portable_lib
25+
from executorch.extension.llm.custom_ops import sdpa_with_kv_cache # noqa
26+
2227
from executorch.extension.pybindings.portable_lib import _load_for_executorch
2328

2429

0 commit comments

Comments
 (0)