Skip to content

Commit cbdb721

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

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

examples/models/llama2/runner/generation.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,11 @@
1919
Message,
2020
Tokenizer,
2121
)
22+
23+
from executorch.extension.pybindings import portable_lib
2224
from executorch.extension.pybindings.portable_lib import _load_for_executorch
2325

26+
from executorch.extension.llm.custom_ops import sdpa_with_kv_cache
2427

2528
class CompletionPrediction(TypedDict, total=False):
2629
generation: str

0 commit comments

Comments
 (0)