small fix (#12623)

intel-analytics · Dec 27, 2024 · 34dbdb8 · 34dbdb8
1 parent 55ce091
commit 34dbdb8
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/python/llm/src/ipex_llm/transformers/models/common.py b/python/llm/src/ipex_llm/transformers/models/common.py
@@ -199,7 +199,7 @@ def prepare_mask(mask, bsz, n_heads, seq_length, kv_length, is_causal, dtype, de
             mask = None
         else:
             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])
     else:
         if seq_length != kv_length and seq_length <= 32:

diff --git a/python/llm/src/ipex_llm/transformers/models/llama.py b/python/llm/src/ipex_llm/transformers/models/llama.py
@@ -68,6 +68,7 @@ def llama_model_forward(
 ) -> Union[Tuple, BaseModelOutputWithPast]:
     # IPEX-LLM OPT start: kv cache and quantize kv cache
     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_quantize_kv = use_quantize_kv_cache(
         self.layers[0].mlp.down_proj, inputs,