ayo/ipex-llm
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

hot-fix mistral kv_cache (#9528)

Browse source
This commit is contained in:
SONG Ge 2023-11-24 14:33:04 +08:00 committed by GitHub
parent 42b7a16bc5
commit 3d24823cda
1 changed files with 5 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
python/llm/src/bigdl/llm/transformers/models/mistral.py
View file

@ -114,6 +114,11 @@ def mistral_attention_forward(
dtype=cache_k.dtype,
device=device)
new_cache_k[:] = cache_k
new_cache_v[:] = cache_v
cache_k = new_cache_k
cache_v = new_cache_v
key_states, value_states = append_kv_cache(cache_k, cache_v, key_states, value_states)
elif use_cache: