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

mistral 4.36 use fp16 sdp (#10704)

Browse source
This commit is contained in:
Yina Chen 2024-04-09 13:50:33 +08:00 committed by GitHub
parent dcb2038aad
commit c7422712fc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 9 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
python/llm/src/ipex_llm/transformers/models/mistral.py
View file

@ -896,6 +896,15 @@ def mistral_attention_forward_4_36_original(
attn_weights = None attn_weights = None
attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
elif use_esimd_sdp(q_len, key_states.shape[2], self.head_dim, query_states):
import linear_fp16_esimd
attn_output = linear_fp16_esimd.sdp_forward(query_states,
key_states,
value_states)
attn_output = attn_output.view(query_states.shape)
attn_weights = None
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
else: else:
attn_output, attn_weights = compute_attn_outputs_weights(query_states, attn_output, attn_weights = compute_attn_outputs_weights(query_states,
key_states, key_states,