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

Fix sdp logic (#10896)

Browse source 
* fix

* fix
This commit is contained in:
Guancheng Fu 2024-04-28 22:02:14 +08:00 committed by GitHub
parent 015d07a58f
commit c9fac8c26b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 25 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1353,10 +1353,31 @@ def llama_attention_forward_4_36_original(
key_states = repeat_kv(key_states, self.num_key_value_groups) key_states = repeat_kv(key_states, self.num_key_value_groups)
value_states = repeat_kv(value_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups)
# otherwise, use native attention # otherwise, use native attention
if query_states.device.type == "xpu":
attn_output, attn_weights = native_sdp(query_states, key_states, value_states, attn_output, attn_weights = native_sdp(query_states, key_states, value_states,
attention_mask, attention_mask,
bsz, q_len, kv_seq_len, bsz, q_len, kv_seq_len,
self.head_dim, self.num_heads, output_attentions) self.head_dim, self.num_heads, output_attentions)
else:
# CPU path
if not output_attentions:
attn_output = torch.nn.functional.scaled_dot_product_attention(
query_states,
key_states,
value_states,
attn_mask=attention_mask,
dropout_p=self.attention_dropout if self.training else 0.0,
# The q_len > 1 is necessary to match with
# AttentionMaskConverter.to_causal_4d that
# does not create a causal mask in case q_len == 1.
is_causal=self.is_causal and attention_mask is None and q_len > 1,
)
else:
attn_output, attn_weights = native_sdp(query_states, key_states, value_states,
attention_mask,
bsz, q_len, kv_seq_len,
self.head_dim,
self.num_heads, output_attentions)
attn_output_size = (bsz, self.num_heads, q_len, self.head_dim) attn_output_size = (bsz, self.num_heads, q_len, self.head_dim)
if attn_output.size() != attn_output_size: if attn_output.size() != attn_output_size:
@ -1778,9 +1799,6 @@ def llama_model_forward_4_36_internal(
if inputs_embeds is None: if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids) inputs_embeds = self.embed_tokens(input_ids)
# IPEX-LLM modifications:
# Disable sdpa for CPU
self._use_sdpa = False
if self._use_flash_attention_2: if self._use_flash_attention_2:
# 2d mask is passed through the layers # 2d mask is passed through the layers
attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) \ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) \