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

small fix (#12623)

Browse source
This commit is contained in:
Yishuo Wang 2024-12-27 10:19:27 +08:00 committed by GitHub
parent 55ce091242
commit 34dbdb8ee3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 2 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -199,7 +199,7 @@ def prepare_mask(mask, bsz, n_heads, seq_length, kv_length, is_causal, dtype, de
mask = None mask = None
else: else:
mask = torch.zeros([1, 1, 1, padding_kv_length], dtype=dtype, device=device) mask = torch.zeros([1, 1, 1, padding_kv_length], dtype=dtype, device=device)
mask[:, :, kv_length:padding_kv_length] = torch.finfo(dtype).min mask[..., kv_length:padding_kv_length] = torch.finfo(dtype).min
mask = mask.expand([bsz, n_heads, seq_length, padding_kv_length]) mask = mask.expand([bsz, n_heads, seq_length, padding_kv_length])
else: else:
if seq_length != kv_length and seq_length <= 32: if seq_length != kv_length and seq_length <= 32:

1
python/llm/src/ipex_llm/transformers/models/llama.py
View file

@ -68,6 +68,7 @@ def llama_model_forward(
) -> Union[Tuple, BaseModelOutputWithPast]: ) -> Union[Tuple, BaseModelOutputWithPast]:
# IPEX-LLM OPT start: kv cache and quantize kv cache # IPEX-LLM OPT start: kv cache and quantize kv cache
inputs = input_ids if input_ids is not None else inputs_embeds inputs = input_ids if input_ids is not None else inputs_embeds
use_cache = use_cache if use_cache is not None else self.config.use_cache
use_cache = True if inputs.device.type == "xpu" else use_cache use_cache = True if inputs.device.type == "xpu" else use_cache
use_quantize_kv = use_quantize_kv_cache( use_quantize_kv = use_quantize_kv_cache(
self.layers[0].mlp.down_proj, inputs, self.layers[0].mlp.down_proj, inputs,