Skip to content

Commit 433bd6e

Browse files
lucylqfacebook-github-bot
authored andcommitted
Add llama sdpa to generation script (#4674)
Summary: After #4024, sdpa is no longer in portable_lib by default, import separately. Differential Revision: D61163753
1 parent 9293b78 commit 433bd6e

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
@@ -17,6 +17,11 @@
1717
from executorch.examples.models.llama2.tokenizer.tiktoken import Tokenizer
1818
from executorch.extension.pybindings.portable_lib import _load_for_executorch
1919

20+
from executorch.extension.pybindings import portable_lib # noqa # usort: skip
21+
22+
# Note: import this after portable_lib
23+
from executorch.extension.llm.custom_ops import sdpa_with_kv_cache # noqa # usort: skip
24+
2025

2126
class CompletionPrediction(TypedDict, total=False):
2227
generation: str

0 commit comments

Comments
 (0)