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

quick fix qwen2 fp8 kv cache (#10135)

Browse source
This commit is contained in:
Yishuo Wang 2024-02-08 17:04:59 +08:00 committed by GitHub
parent 39d90839aa
commit 4d33aac7f9
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/bigdl/llm/transformers/models/qwen2.py
View file

@ -167,6 +167,8 @@ def qwen2_attention_forward_quantized(
if q_len != 1: if q_len != 1:
key, value = restore_fp8_kv_cache(key_states, value_states, query_states.dtype) key, value = restore_fp8_kv_cache(key_states, value_states, query_states.dtype)
key = repeat_kv(key, self.num_key_value_groups)
value = repeat_kv(value, self.num_key_value_groups)
attn_weights = torch.matmul(query_states, key.transpose(2, 3)) attn_weights = torch.matmul(query_states, key.transpose(2, 3))
else: else:
import linear_q4_0 import linear_q4_0