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

use new fp32 softmax kernel (#11776)

Browse source
This commit is contained in:
Yishuo Wang 2024-08-13 14:48:11 +08:00 committed by GitHub
parent 23d3acdc77
commit aa861df066
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 6 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
python/llm/src/ipex_llm/transformers/models/minicpmv.py
View file

@ -42,8 +42,9 @@ def siglip_attention_forward(
if attention_mask is not None:
attn_weights = attn_weights + attention_mask
# upcast attention to fp32
attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)
import xe_addons
xe_addons.attn_softmax_inplaced(attn_weights)
attn_weights = torch.nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
attn_output = torch.matmul(attn_weights, value_states)

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

@ -184,9 +184,9 @@ def attention_forward(
if attention_mask is not None:
attn_weights = attn_weights + attention_mask
# upcast attention to fp32
attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1,
dtype=torch.float32).to(value_states.dtype)
import xe_addons
xe_addons.attn_softmax_inplaced(attn_weights)
attn_weights = torch.nn.functional.dropout(attn_weights, p=self.attention_dropout,
training=self.training)
attn_output = torch.matmul(attn_weights, value_states)