diff --git a/python/llm/src/ipex_llm/transformers/models/baichuan.py b/python/llm/src/ipex_llm/transformers/models/baichuan.py index 33764ba466c..a78e5f8e131 100644 --- a/python/llm/src/ipex_llm/transformers/models/baichuan.py +++ b/python/llm/src/ipex_llm/transformers/models/baichuan.py @@ -29,7 +29,7 @@ should_use_compresskv from ipex_llm.transformers.models.utils import update_past_key_value from ipex_llm.transformers.models.utils import should_use_fuse_rope -from ipex_llm.transformers.models.utils import use_flash_attention, use_sdp +from ipex_llm.transformers.models.utils import use_sdp from ipex_llm.transformers.models.utils import apply_rotary_pos_emb, SILU from ipex_llm.transformers.models.utils import mlp_fusion_check from ipex_llm.transformers.models.utils import is_enough_kv_cache_room_4_36 @@ -301,16 +301,10 @@ def baichuan_attention_forward_7b( # IPEX-LLM OPT: sdp attn_weights = None - if use_flash_attention(query_states, key_states, attention_mask): - attn_output = F.scaled_dot_product_attention(query_states.to(dtype=torch.float16), - key_states.to(dtype=torch.float16), - value_states.to(dtype=torch.float16), - is_causal=True).to(hidden_states.dtype) - else: - attn_output = scaled_dot_product_attention( - query_states, key_states, value_states, - attention_mask, q_len == kv_seq_len - ) + attn_output = scaled_dot_product_attention( + query_states, key_states, value_states, + attention_mask, q_len == kv_seq_len + ) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) diff --git a/python/llm/src/ipex_llm/transformers/models/chatglm.py b/python/llm/src/ipex_llm/transformers/models/chatglm.py index 77e2ae4491d..34241d89d97 100644 --- a/python/llm/src/ipex_llm/transformers/models/chatglm.py +++ b/python/llm/src/ipex_llm/transformers/models/chatglm.py @@ -23,7 +23,7 @@ import torch.nn.functional as F from typing import Optional, Tuple from ipex_llm.transformers.models.utils import init_kv_cache, extend_kv_cache, append_kv_cache -from ipex_llm.transformers.models.utils import use_flash_attention, use_sdp +from ipex_llm.transformers.models.utils import use_sdp def rotate_half(x): @@ -41,7 +41,7 @@ def apply_rotary_pos_emb_index(q, k, cos, sin, position_id): def glm_sdpa(query, key, value, attention_mask=None, is_causal=False): - if use_flash_attention(query, key, attention_mask) or query.device.type == 'cpu': + if query.device.type == 'cpu': context_layer = F.scaled_dot_product_attention(query.to(key.dtype), key, value, diff --git a/python/llm/src/ipex_llm/transformers/models/qwen.py b/python/llm/src/ipex_llm/transformers/models/qwen.py index dcebdee33c6..590867c85ff 100644 --- a/python/llm/src/ipex_llm/transformers/models/qwen.py +++ b/python/llm/src/ipex_llm/transformers/models/qwen.py @@ -33,7 +33,6 @@ from ipex_llm.transformers.models.utils import use_quantize_kv_cache from ipex_llm.transformers.models.utils import rotate_half, SILU from ipex_llm.transformers.models.utils import mlp_fusion_check -from ipex_llm.transformers.models.utils import use_flash_attention from ipex_llm.utils.common import invalidInputError from transformers.modeling_outputs import BaseModelOutputWithPast @@ -116,33 +115,28 @@ def qwen_attention_forward( past_key_value = (key_states.transpose(1, 2), value_states.transpose(1, 2)) if use_cache else None - # IPEX-LLM OPT: sdp + # IPEX-LLM OPT: sdpa attn_weights = None - if use_flash_attention(query_states, key_states, attention_mask): - attn_output = F.scaled_dot_product_attention(query_states.to(dtype=torch.float16), - key_states.to(dtype=torch.float16), - value_states.to(dtype=torch.float16), - is_causal=True).to(hidden_states.dtype) + + if q_len > 1 and q_len != kv_seq_len: + causal_mask = torch.tril( + torch.ones((kv_seq_len, kv_seq_len), dtype=torch.bool, device=query_states.device) + ).view(1, 1, kv_seq_len, kv_seq_len) + causal_mask = causal_mask[ + :, :, kv_seq_len - q_len:kv_seq_len, :kv_seq_len + ] + attention_mask = torch.zeros(causal_mask.shape, dtype=query_states.dtype, + device=query_states.device) + attention_mask.masked_fill_(causal_mask.logical_not(), + torch.finfo(attention_mask.dtype).min) + attention_mask = attention_mask.expand([bsz, -1, -1, -1]) else: - if q_len > 1 and q_len != kv_seq_len: - causal_mask = torch.tril( - torch.ones((kv_seq_len, kv_seq_len), dtype=torch.bool, device=query_states.device) - ).view(1, 1, kv_seq_len, kv_seq_len) - causal_mask = causal_mask[ - :, :, kv_seq_len - q_len:kv_seq_len, :kv_seq_len - ] - attention_mask = torch.zeros(causal_mask.shape, dtype=query_states.dtype, - device=query_states.device) - attention_mask.masked_fill_(causal_mask.logical_not(), - torch.finfo(attention_mask.dtype).min) - attention_mask = attention_mask.expand([bsz, -1, -1, -1]) - else: - attention_mask = None + attention_mask = None - attn_output = scaled_dot_product_attention( - query_states, key_states, value_states, - attention_mask, q_len == kv_seq_len - ) + attn_output = scaled_dot_product_attention( + query_states, key_states, value_states, + attention_mask, q_len == kv_seq_len + ) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(bsz, q_len, self.hidden_size) @@ -219,31 +213,25 @@ def qwen_attention_forward_registered( past_key_value = (key_states.transpose(1, 2), value_states.transpose(1, 2)) if use_cache else None - # IPEX-LLM OPT: sdp + # IPEX-LLM OPT: sdpa attn_weights = None - if use_flash_attention(query_states, key_states, attention_mask): - attn_output = F.scaled_dot_product_attention(query_states.to(dtype=torch.float16), - key_states.to(dtype=torch.float16), - value_states.to(dtype=torch.float16), - is_causal=True).to(hidden_states.dtype) + if q_len > 1 and q_len != kv_seq_len: + causal_mask = registered_causal_mask[ + :, :, kv_seq_len - q_len:kv_seq_len, :kv_seq_len + ] + attention_mask = torch.zeros(causal_mask.shape, dtype=query_states.dtype, + device=query_states.device) + attention_mask.masked_fill_(causal_mask.logical_not(), + torch.finfo(attention_mask.dtype).min) + attention_mask = attention_mask.expand([bsz, -1, -1, -1]) else: - if q_len > 1 and q_len != kv_seq_len: - causal_mask = registered_causal_mask[ - :, :, kv_seq_len - q_len:kv_seq_len, :kv_seq_len - ] - attention_mask = torch.zeros(causal_mask.shape, dtype=query_states.dtype, - device=query_states.device) - attention_mask.masked_fill_(causal_mask.logical_not(), - torch.finfo(attention_mask.dtype).min) - attention_mask = attention_mask.expand([bsz, -1, -1, -1]) - else: - attention_mask = None + attention_mask = None - attn_output = scaled_dot_product_attention( - query_states, key_states, value_states, - attention_mask, q_len == kv_seq_len - ) + attn_output = scaled_dot_product_attention( + query_states, key_states, value_states, + attention_mask, q_len == kv_seq_len + ) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(bsz, q_len, self.hidden_size) diff --git a/python/llm/src/ipex_llm/transformers/models/qwen2.py b/python/llm/src/ipex_llm/transformers/models/qwen2.py index 324ae2aa85e..62f48e2d012 100644 --- a/python/llm/src/ipex_llm/transformers/models/qwen2.py +++ b/python/llm/src/ipex_llm/transformers/models/qwen2.py @@ -38,12 +38,10 @@ # import os -import math from typing import Optional, Tuple, Union, List import torch from torch.nn import CrossEntropyLoss -from torch.nn.functional import scaled_dot_product_attention as sdpa from ipex_llm.transformers.models.common import merge_qkv_base from ipex_llm.transformers.models.common import scaled_dot_product_attention @@ -51,13 +49,12 @@ from ipex_llm.transformers.models.utils import should_use_fuse_rope from ipex_llm.transformers.models.utils import use_quantize_kv_cache, \ should_use_compresskv, is_enough_kv_cache_room_4_36 -from ipex_llm.transformers.models.utils import use_flash_attention from ipex_llm.transformers.kv import DynamicFp8Cache, DynamicNormalCache, \ DynamicCompressCache, DynamicCompressFp8Cache from ipex_llm.utils.common import invalidInputError from transformers.models.qwen2.modeling_qwen2 import Qwen2Attention, Qwen2MLP -from transformers.models.qwen2.modeling_qwen2 import apply_rotary_pos_emb, repeat_kv +from transformers.models.qwen2.modeling_qwen2 import apply_rotary_pos_emb from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.cache_utils import Cache from transformers import logging @@ -580,21 +577,10 @@ def qwen2_attention_forward( self.layer_idx, None) attn_weights = None - if use_flash_attention(query_states, key_states, attention_mask): - if attention_mask is not None: - attention_mask = attention_mask[:, :, :, :kv_seq_len] - # 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) - attn_output = sdpa(query_states.to(device, dtype=torch.float16), - key_states.to(device, dtype=torch.float16), - value_states.to(device, dtype=torch.float16), - is_causal=True).to(hidden_states.dtype) - else: - attn_output = scaled_dot_product_attention( - query_states, key_states, value_states, - attention_mask, q_len == kv_seq_len - ) + attn_output = scaled_dot_product_attention( + query_states, key_states, value_states, + attention_mask, q_len == kv_seq_len + ) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)