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

use sdp fp8 causal kernel (#11023)

Browse source
This commit is contained in:
Yishuo Wang 2024-05-15 10:22:35 +08:00 committed by GitHub
parent c34f85e7d0
commit fad1dbaf60
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 24 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
python/llm/src/ipex_llm/transformers/models/phi3.py
View file

@ -137,22 +137,23 @@ def attention_forward(
key_states, value_states = past_key_value.update(key_states, value_states,
self.layer_idx, None)
if (isinstance(past_key_value, DynamicFp8Cache) and
use_sdp_fp8(q_len, kv_seq_len, query_states)):
import linear_q4_0
attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, attention_mask)
elif (isinstance(past_key_value, DynamicNormalCache) and
use_sdp(q_len, kv_seq_len, self.head_dim, query_states)):
if use_sdp(q_len, kv_seq_len, self.head_dim, query_states):
import linear_q4_0
if isinstance(past_key_value, DynamicFp8Cache):
attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states,
attention_mask)
else:
attn_output = linear_q4_0.sdp(query_states, key_states, value_states, attention_mask)
elif use_sdp_causal(q_len, kv_seq_len, query_states, self.training):
import linear_q4_0
if isinstance(past_key_value, DynamicFp8Cache):
attn_output = linear_q4_0.sdp_fp8_causal(query_states, key_states, value_states)
else:
attn_output = linear_q4_0.sdp_causal(query_states, key_states, value_states)
else:
if isinstance(past_key_value, DynamicFp8Cache):
key_states, value_states = restore_fp8_kv_cache(key_states, value_states,
query_states.dtype)
if use_sdp_causal(q_len, kv_seq_len, query_states, self.training):
import linear_q4_0
attn_output = linear_q4_0.sdp_causal(query_states, key_states, value_states)
else:
# repeat k/v heads if n_kv_heads < n_heads
key_states = repeat_kv(key_states, self.num_key_value_groups)
value_states = repeat_kv(value_states, self.num_key_value_groups)