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

add comment

Browse source
This commit is contained in:
Huang, Xinshengzi 2024-08-22 15:14:47 +08:00
parent a2be3d7501
commit eb1e65f8a9
1 changed files with 2 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
python/llm/src/ipex_llm/transformers/models/baichuan.py
View file

@ -246,6 +246,7 @@ def baichuan_attention_forward_7b(
bsz, q_len, _ = hidden_states.size() bsz, q_len, _ = hidden_states.size()
device = hidden_states.device device = hidden_states.device
# [CompressKV]
use_compresskv = isinstance(past_key_value, DynamicCompressCache) use_compresskv = isinstance(past_key_value, DynamicCompressCache)
qkv = self.W_pack(hidden_states) qkv = self.W_pack(hidden_states)
@ -257,6 +258,7 @@ def baichuan_attention_forward_7b(
kv_seq_len = key_states.shape[2] kv_seq_len = key_states.shape[2]
if past_key_value is not None: if past_key_value is not None:
# [CompressKV]
if use_compresskv: if use_compresskv:
kv_seq_len += past_key_value.get_usable_length(kv_seq_len, kv_seq_len += past_key_value.get_usable_length(kv_seq_len,
self.layer_idx) self.layer_idx)