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

[NPU] Qwen prefill attn_mask type hotfix (#12395)

Browse source 
* qwen prefill attn_mask type fp16

* update
This commit is contained in:
Yina Chen 2024-11-13 11:51:34 +02:00 committed by GitHub
parent 9220babaab
commit d6d63d6b84
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
python/llm/src/ipex_llm/transformers/npu_models/mp_models_base.py
View file

@ -247,7 +247,8 @@ class LLMBaseNNFactory(NNFactory):
attn_weight = self.matmul(query_states, key_states, False, True) / (
math.sqrt(head_dim)
)
attention_mask = self.convert_to_fp16(attention_mask)
if mode != "prefill":
attention_mask = self.convert_to_fp16(attention_mask)
attn_weight = self.eltwise_add(attn_weight, attention_mask)
attn_weight = self.convert_to_fp32(attn_weight)
attn_weight = self.softmax(attn_weight, -1)