diff --git a/python/llm/src/ipex_llm/transformers/models/baichuan2.py b/python/llm/src/ipex_llm/transformers/models/baichuan2.py index 550035c3..e0c76a5d 100644 --- a/python/llm/src/ipex_llm/transformers/models/baichuan2.py +++ b/python/llm/src/ipex_llm/transformers/models/baichuan2.py @@ -49,6 +49,21 @@ import os KV_CACHE_ALLOC_BLOCK_LENGTH = int(os.environ.get("KV_CACHE_ALLOC_BLOCK_LENGTH", 256)) +def should_split_qkv_tensor(query_states, bsz, num_heads, q_len, kv_seq_len, output_attentions): + if not output_attentions: + if os.environ.get("IPEX_LLM_SPLIT_QKV", None) is not None: + return os.environ.get("IPEX_LLM_SPLIT_QKV", None) == "1" + elif query_states.dtype == torch.float16 and \ + query_states.shape[2] >= 5400: + # split tensor for memory block limitation + # support fp16 and set input length threshold at 5400 for now + return True + elif query_states.element_size()*bsz*num_heads*q_len*kv_seq_len >= 4*1024**3: + # attn_weight size larger than memory block limitation 4GB + return True + return False + + def baichuan_13b_rms_norm_forward(self, hidden_states): if hidden_states.device.type == "xpu" and not (self.training and hidden_states.requires_grad): import linear_q4_0 @@ -159,13 +174,18 @@ def baichuan_attention_forward_7b_quantized( if query_states.size(2) != 1 or device.type != 'xpu': key_states, value_states = restore_fp8_kv_cache(key_states, value_states, query_states.dtype) - attn_output = torch.matmul(query_states * scaling_factor, key_states.transpose(-2, -1)) + if should_split_qkv_tensor(query_states, bsz, self.num_heads, + q_len, kv_seq_len, output_attentions): + attn_output, attn_weights = native_sdp_split_qkv_tensor(query_states, key_states, + value_states, attention_mask) + else: + attn_output = torch.matmul(query_states * scaling_factor, key_states.transpose(-2, -1)) - if attention_mask is not None: - attn_output += attention_mask - attn_output = torch.softmax(attn_output, -1) - attn_output = attn_output.to(hidden_states.dtype) - attn_output = torch.matmul(attn_output, value_states) + if attention_mask is not None: + attn_output += attention_mask + attn_output = torch.softmax(attn_output, -1) + attn_output = attn_output.to(hidden_states.dtype) + attn_output = torch.matmul(attn_output, value_states) else: import linear_q4_0 attn_output = linear_q4_0.sdp_fp8(query_states, key_states, value_states, @@ -287,13 +307,20 @@ def baichuan_attention_forward_7b_origin( if attention_mask is not None: if attention_mask.dtype == torch.bool: attention_mask.masked_fill_(attention_mask.logical_not(), float("-inf")) - - scaling_factor = 1 / math.sqrt(query_states.size(-1)) - attn_output = torch.matmul(query_states * scaling_factor, key_states.transpose(-2, -1)) - if attention_mask is not None: - attn_output += attention_mask - attn_output = torch.softmax(attn_output, -1) - attn_output = torch.matmul(attn_output, value_states) + if should_split_qkv_tensor(query_states, bsz, self.num_heads, + q_len, kv_seq_len, output_attentions): + attn_output, attn_weights = native_sdp_split_qkv_tensor(query_states, + key_states, + value_states, + attention_mask) + else: + scaling_factor = 1 / math.sqrt(query_states.size(-1)) + attn_output = torch.matmul(query_states * scaling_factor, + key_states.transpose(-2, -1)) + if attention_mask is not None: + attn_output += attention_mask + attn_output = torch.softmax(attn_output, -1) + attn_output = torch.matmul(attn_output, value_states) attn_output = attn_output.transpose(1, 2) attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) @@ -622,3 +649,21 @@ def baichuan_13b_get_alibi_mask(self, tensor, seq_length_with_past): : self.n_head, :seq_length_with_past, :seq_length_with_past ] return mask + + +def native_sdp_split_qkv_tensor(query, key, value, attention_mask): + block_size = 8 + query_split = torch.split(query, block_size, dim=1) + key_split = torch.split(key.transpose(-2, -1), block_size, dim=1) + value_split = torch.split(value, block_size, dim=1) + attn_outputs = [] + scaling_factor = 1 / math.sqrt(query.size(-1)) + for q, k, v in zip(query_split, key_split, value_split): + attn_output_split = torch.matmul(q * scaling_factor, k) + if attention_mask is not None: + attn_output_split += attention_mask + attn_output_split = torch.softmax(attn_output_split, -1) + attn_output_split = torch.matmul(attn_output_split, v) + attn_outputs.append(attn_output_split) + attn_output = torch.cat(attn_outputs, dim=1) + return attn_output, None diff --git a/python/llm/src/ipex_llm/transformers/models/mistral.py b/python/llm/src/ipex_llm/transformers/models/mistral.py index 10791a06..ef1971f2 100644 --- a/python/llm/src/ipex_llm/transformers/models/mistral.py +++ b/python/llm/src/ipex_llm/transformers/models/mistral.py @@ -89,6 +89,21 @@ def should_use_fuse_rope(self, hidden_states, position_ids): return use_fuse_rope +def should_split_qkv_tensor(query_states, bsz, num_heads, q_len, kv_seq_len, output_attentions): + if not output_attentions: + if os.environ.get("IPEX_LLM_SPLIT_QKV", None) is not None: + return os.environ.get("IPEX_LLM_SPLIT_QKV", None) == "1" + elif query_states.dtype == torch.float16 and \ + query_states.shape[2] >= 6300: + # split tensor for memory block limitation + # support fp16 and set input length threshold at 6300 for now + return True + elif query_states.element_size()*bsz*num_heads*q_len*kv_seq_len >= 4*1024**3: + # attn_weight size larger than memory block limitation 4GB + return True + return False + + def compute_attn_outputs_weights(query_states, key_states, value_states, bsz, q_len, kv_seq_len, num_heads, head_dim, hidden_size, attention_mask): attn_weights = torch.matmul( @@ -112,9 +127,14 @@ def compute_attn_outputs_weights(query_states, key_states, value_states, bsz, q_ attn_weights = attn_weights + attention_mask - # upcast attention to fp32 - attn_weights = nn.functional.\ - softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + if kv_seq_len >= 2048 or bsz >= 64: + # for memory considerations, do not upcast attention to fp32 + # for long sequences or large batches + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + else: + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, + dtype=torch.float32).to(query_states.dtype) attn_output = torch.matmul(attn_weights, value_states.to(query_states.dtype)) if attn_output.size() != (bsz, num_heads, q_len, head_dim): @@ -130,6 +150,45 @@ def compute_attn_outputs_weights(query_states, key_states, value_states, bsz, q_ return attn_output, attn_weights +def compute_attn_outputs_weights_split_tensor(query_states, key_states, value_states, + bsz, q_len, kv_seq_len, num_heads, head_dim, + hidden_size, attention_mask): + block_size = 8 + query_split = torch.split(query_states.to(key_states.dtype), block_size, dim=1) + key_split = torch.split(key_states.transpose(2, 3), block_size, dim=1) + value_split = torch.split(value_states.to(query_states.dtype), block_size, dim=1) + attn_outputs = [] + for q, k, v in zip(query_split, key_split, value_split): + attn_weights_split = torch.matmul(q, k) / math.sqrt(head_dim) + block_actual_size = attn_weights_split.size(1) + attn_weights_split_size = (bsz, block_actual_size, q_len, kv_seq_len) + if attn_weights_split.size() != attn_weights_split_size: + invalidInputError(False, + f"Splitted attention weights should be of size " + f"{attn_weights_split_size}, but is {attn_weights_split.size()}") + + if attention_mask is not None: + attn_mask_size = (bsz, 1, q_len, kv_seq_len) + if attention_mask.size() != attn_mask_size: + invalidInputError(False, + f"Attention mask should be of size {attn_mask_size}, " + f"but is {attention_mask.size()}") + attn_weights_split = attn_weights_split + attention_mask + attn_weights_split = nn.functional.softmax(attn_weights_split, dim=-1) + attn_outputs.append(torch.matmul(attn_weights_split, v)) + attn_output = torch.cat(attn_outputs, dim=1) + if attn_output.size() != (bsz, num_heads, q_len, head_dim): + invalidInputError( + False, + f"`attn_output` should be of size {(bsz, num_heads, q_len, head_dim)}," + f" but is {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, hidden_size) + return attn_output, None + + def mistral_model_forward_4_36( self, input_ids: torch.LongTensor = None, @@ -272,30 +331,58 @@ def mistral_attention_forward_quantized( dtype=attention_dtype) kv_seq_len = key_states.shape[-2] if past_key_value is None: - attn_weights = torch.matmul(query_states.to(key_states.dtype), - key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + if should_split_qkv_tensor(query_states, bsz, self.num_heads, + q_len, kv_seq_len, output_attentions): + block_size = 8 + query_split = torch.split(query_states.to(key_states.dtype), block_size, dim=1) + key_split = torch.split(key_states.transpose(2, 3), block_size, dim=1) + value_split = torch.split(value_states.to(query_states.dtype), block_size, dim=1) + attn_outputs = [] + for q, k, v in zip(query_split, key_split, value_split): + attn_weights_split = torch.matmul(q, k) / math.sqrt(self.head_dim) + block_actual_size = attn_weights_split.size(1) + attn_weights_split_size = (bsz, block_actual_size, q_len, kv_seq_len) + if attn_weights_split.size() != attn_weights_split_size: + invalidInputError(False, + f"Splitted attention weights should be of size " + f"{attn_weights_split_size}, " + f"but is {attn_weights_split.size()}") - if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): - invalidInputError( - False, - f"Attention weights should be of size " - f"{(bsz, self.num_heads, q_len, kv_seq_len)}, but is" - f" {attn_weights.size()}" - ) + if attention_mask is not None: + attn_mask_size = (bsz, 1, q_len, kv_seq_len) + if attention_mask.size() != attn_mask_size: + invalidInputError(False, + f"Attention mask should be of size {attn_mask_size}, " + f"but is {attention_mask.size()}") + attn_weights_split = attn_weights_split + attention_mask + attn_weights_split = nn.functional.softmax(attn_weights_split, dim=-1) + attn_outputs.append(torch.matmul(attn_weights_split, v)) + attn_output = torch.cat(attn_outputs, dim=1) + else: + attn_weights = torch.matmul(query_states.to(key_states.dtype), + key_states.transpose(2, 3)) / math.sqrt(self.head_dim) - if attention_mask is not None: - if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): invalidInputError( False, - f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}," - f" but is {attention_mask.size()}" + f"Attention weights should be of size " + f"{(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" ) - attn_weights = attn_weights + attention_mask - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dim=-1, - dtype=torch.float32).to(query_states.dtype) - attn_output = torch.matmul(attn_weights, value_states) + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + invalidInputError( + False, + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}," + f" but is {attention_mask.size()}" + ) + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, + dtype=torch.float32).to(query_states.dtype) + attn_output = torch.matmul(attn_weights, value_states) if use_cache: k_cache, v_cache = init_fp8_kv_cache( bsz, self.num_heads, kv_seq_len, self.head_dim, @@ -518,12 +605,29 @@ def mistral_attention_forward_original( dtype=attention_dtype) value_states = repeat_kv(value_states, self.num_key_value_groups).to(device, dtype=attention_dtype) - attn_output, attn_weights = compute_attn_outputs_weights(query_states, - key_states, - value_states, - bsz, q_len, kv_seq_len, - self.num_heads, self.head_dim, - self.hidden_size, attention_mask) + if should_split_qkv_tensor(query_states, bsz, self.num_heads, + q_len, kv_seq_len, output_attentions): + attn_output, attn_weights = compute_attn_outputs_weights_split_tensor(query_states, + key_states, + value_states, + bsz, + q_len, + kv_seq_len, + self.num_heads, + self.head_dim, + self.hidden_size, + attention_mask) + else: + attn_output, attn_weights = compute_attn_outputs_weights(query_states, + key_states, + value_states, + bsz, + q_len, + kv_seq_len, + self.num_heads, + self.head_dim, + self.hidden_size, + attention_mask) attn_output = self.o_proj(attn_output) @@ -653,30 +757,63 @@ def mistral_attention_forward_4_36_quantized( dtype=attention_dtype) kv_seq_len = key_states.shape[-2] if len(past_key_value.key_cache) <= self.layer_idx: - attn_weights = torch.matmul(query_states.to(key_states.dtype), - key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + if should_split_qkv_tensor(query_states, bsz, self.num_heads, + q_len, kv_seq_len, output_attentions): + block_size = 8 + query_split = torch.split(query_states.to(key_states.dtype), block_size, dim=1) + key_split = torch.split(key_states.transpose(2, 3), block_size, dim=1) + value_split = torch.split(value_states.to(query_states.dtype), block_size, dim=1) + attn_outputs = [] + for q, k, v in zip(query_split, key_split, value_split): + attn_weights_split = torch.matmul(q, k) / math.sqrt(self.head_dim) + block_actual_size = attn_weights_split.size(1) + attn_weights_split_size = (bsz, block_actual_size, q_len, kv_seq_len) + if attn_weights_split.size() != attn_weights_split_size: + invalidInputError(False, + f"Splitted attention weights should be of size " + f"{attn_weights_split_size}, " + f"but is {attn_weights_split.size()}") - if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): - invalidInputError( - False, - f"Attention weights should be of size " - f"{(bsz, self.num_heads, q_len, kv_seq_len)}, but is" - f" {attn_weights.size()}" - ) + if attention_mask is not None: + attn_mask_size = (bsz, 1, q_len, kv_seq_len) + if attention_mask.size() != attn_mask_size: + invalidInputError(False, + f"Attention mask should be of size {attn_mask_size}, " + f"but is {attention_mask.size()}") + attn_weights_split = attn_weights_split + attention_mask + attn_weights_split = nn.functional.softmax(attn_weights_split, dim=-1) + attn_outputs.append(torch.matmul(attn_weights_split, v)) + attn_output = torch.cat(attn_outputs, dim=1) + else: + attn_weights = torch.matmul(query_states.to(key_states.dtype), + key_states.transpose(2, 3)) / math.sqrt(self.head_dim) - if attention_mask is not None: - if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): invalidInputError( False, - f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}," - f" but is {attention_mask.size()}" + f"Attention weights should be of size " + f"{(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" ) - attn_weights = attn_weights + attention_mask - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dim=-1, - dtype=torch.float32).to(query_states.dtype) - attn_output = torch.matmul(attn_weights, value_states) + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + invalidInputError( + False, + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}," + f" but is {attention_mask.size()}" + ) + attn_weights = attn_weights + attention_mask + + if kv_seq_len >= 2048 or bsz >= 64: + # for memory considerations, do not upcast attention to fp32 + # for long sequences or large batches + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + else: + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, + dtype=torch.float32).to(query_states.dtype) + attn_output = torch.matmul(attn_weights, value_states) if use_cache: cache_kwargs = None key_states, value_states = past_key_value.update(key_states, value_states, @@ -909,14 +1046,29 @@ def mistral_attention_forward_4_36_original( dtype=attention_dtype) value_states = repeat_kv(value_states, self.num_key_value_groups).to(device, dtype=attention_dtype) - attn_output, attn_weights = compute_attn_outputs_weights(query_states, - key_states, - value_states, - bsz, q_len, kv_seq_len, - self.num_heads, - self.head_dim, - self.hidden_size, - attention_mask) + if should_split_qkv_tensor(query_states, bsz, self.num_heads, + q_len, kv_seq_len, output_attentions): + attn_output, attn_weights = compute_attn_outputs_weights_split_tensor(query_states, + key_states, + value_states, + bsz, + q_len, + kv_seq_len, + self.num_heads, + self.head_dim, + self.hidden_size, + attention_mask) + else: + attn_output, attn_weights = compute_attn_outputs_weights(query_states, + key_states, + value_states, + bsz, + q_len, + kv_seq_len, + self.num_heads, + self.head_dim, + self.hidden_size, + attention_mask) attn_output = self.o_proj(attn_output) diff --git a/python/llm/src/ipex_llm/transformers/models/qwen2.py b/python/llm/src/ipex_llm/transformers/models/qwen2.py index 7494e617..13523092 100644 --- a/python/llm/src/ipex_llm/transformers/models/qwen2.py +++ b/python/llm/src/ipex_llm/transformers/models/qwen2.py @@ -74,6 +74,21 @@ import os KV_CACHE_ALLOC_BLOCK_LENGTH = int(os.environ.get("KV_CACHE_ALLOC_BLOCK_LENGTH", 256)) +def should_split_qkv_tensor(query_states, bsz, num_heads, q_len, kv_seq_len, output_attentions): + if not output_attentions: + if os.environ.get("IPEX_LLM_SPLIT_QKV", None) is not None: + return os.environ.get("IPEX_LLM_SPLIT_QKV", None) == "1" + elif query_states.dtype == torch.float16 and \ + query_states.shape[2] >= 5000: + # split tensor for memory block limitation + # support fp16 and set input length threshold at 5000 for now + return True + elif query_states.element_size()*bsz*num_heads*q_len*kv_seq_len >= 4*1024**3: + # attn_weight size larger than memory block limitation 4GB + return True + return False + + def should_use_fuse_rope(self, query_states, position_ids): use_fuse_rope = query_states.device.type == "xpu" use_fuse_rope = use_fuse_rope and not (self.training and query_states.requires_grad) @@ -370,28 +385,43 @@ def qwen2_attention_forward_quantized( key, value = restore_fp8_kv_cache(key_states, value_states, query_states.dtype) key = repeat_kv(key, self.num_key_value_groups) value = repeat_kv(value, self.num_key_value_groups) - attn_weights = torch.matmul(query_states, key.transpose(2, 3)) - attn_weights = attn_weights / math.sqrt(self.head_dim) + if should_split_qkv_tensor(query_states, bsz, self.num_heads, + q_len, kv_seq_len, output_attentions): + attn_output, attn_weights = native_sdp_split_qkv_tensor(query_states, key, + value, attention_mask, + bsz, q_len, kv_seq_len, + self.head_dim, self.num_heads, + self.attention_dropout, + self.training) + else: + attn_weights = torch.matmul(query_states, key.transpose(2, 3)) + attn_weights = attn_weights / math.sqrt(self.head_dim) - invalidInputError(attn_weights.size() == (bsz, self.num_heads, q_len, kv_seq_len), - ("Attention weights should be of size " - f"{(bsz, self.num_heads, q_len, kv_seq_len)}," - "but is {attn_weights.size()}")) + invalidInputError(attn_weights.size() == (bsz, self.num_heads, q_len, kv_seq_len), + ("Attention weights should be of size " + f"{(bsz, self.num_heads, q_len, kv_seq_len)}," + "but is {attn_weights.size()}")) - if attention_mask is not None: - invalidInputError(attention_mask.size() == (bsz, 1, q_len, kv_seq_len), - (f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}" - f" but is {attention_mask.size()}")) + if attention_mask is not None: + invalidInputError(attention_mask.size() == (bsz, 1, q_len, kv_seq_len), + (f"Attention mask should be of size " + f"{(bsz, 1, q_len, kv_seq_len)}," + f" but is {attention_mask.size()}")) - attn_weights = attn_weights + attention_mask + attn_weights = attn_weights + attention_mask - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dim=-1, - dtype=torch.float32).to(query_states.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, - training=self.training) + if kv_seq_len >= 2048 or bsz >= 64: + # for memory considerations, do not upcast attention to fp32 + # for long sequences or large batches + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + else: + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, + dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, + training=self.training) - attn_output = torch.matmul(attn_weights, value) + attn_output = torch.matmul(attn_weights, value) invalidInputError(attn_output.size() == (bsz, self.num_heads, q_len, self.head_dim), "`attn_output` should be of size " @@ -543,28 +573,43 @@ def qwen2_attention_forward_origin( attn_output = attn_output.view(query_states.shape) attn_weights = None else: - attn_weights = torch.matmul(query_states, - key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + if should_split_qkv_tensor(query_states, bsz, self.num_heads, + q_len, kv_seq_len, output_attentions): + attn_output, attn_weights = native_sdp_split_qkv_tensor(query_states, key_states, + value_states, attention_mask, + bsz, q_len, kv_seq_len, + self.head_dim, self.num_heads, + self.attention_dropout, + self.training) + else: + attn_weights = torch.matmul(query_states, + key_states.transpose(2, 3)) / math.sqrt(self.head_dim) - invalidInputError(attn_weights.size() == (bsz, self.num_heads, q_len, kv_seq_len), - ("Attention weights should be of size " - f"{(bsz, self.num_heads, q_len, kv_seq_len)}," - "but is {attn_weights.size()}")) + invalidInputError(attn_weights.size() == (bsz, self.num_heads, q_len, kv_seq_len), + ("Attention weights should be of size " + f"{(bsz, self.num_heads, q_len, kv_seq_len)}," + "but is {attn_weights.size()}")) - if attention_mask is not None: - invalidInputError(attention_mask.size() == (bsz, 1, q_len, kv_seq_len), - (f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}" - f" but is {attention_mask.size()}")) + if attention_mask is not None: + invalidInputError(attention_mask.size() == (bsz, 1, q_len, kv_seq_len), + (f"Attention mask should be of size " + f"{(bsz, 1, q_len, kv_seq_len)}," + f" but is {attention_mask.size()}")) - attn_weights = attn_weights + attention_mask + attn_weights = attn_weights + attention_mask - # upcast attention to fp32 - attn_weights = \ - nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) - attn_weights = nn.functional.dropout(attn_weights, - p=self.attention_dropout, - training=self.training) - attn_output = torch.matmul(attn_weights, value_states) + if kv_seq_len >= 2048 or bsz >= 64: + # for memory considerations, do not upcast attention to fp32 + # for long sequences or large batches + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + else: + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, + dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, + p=self.attention_dropout, + training=self.training) + attn_output = torch.matmul(attn_weights, value_states) invalidInputError(attn_output.size() == (bsz, self.num_heads, q_len, self.head_dim), "`attn_output` should be of size " @@ -725,3 +770,36 @@ def qwen2_sdpa_attention_forward( attn_output = self.o_proj(attn_output) return attn_output, None, past_key_value + + +def native_sdp_split_qkv_tensor(query, key, value, attention_mask, + bsz, q_len, kv_seq_len, head_dim, num_heads, + attention_dropout, training): + block_size = 8 + query_split = torch.split(query, block_size, dim=1) + key_split = torch.split(key.transpose(2, 3), block_size, dim=1) + value_split = torch.split(value, block_size, dim=1) + attn_outputs = [] + for q, k, v in zip(query_split, key_split, value_split): + attn_weights_split = torch.matmul(q, k) / math.sqrt(head_dim) + block_actual_size = attn_weights_split.size(1) + attn_weights_split_size = (bsz, block_actual_size, q_len, kv_seq_len) + if attn_weights_split.size() != attn_weights_split_size: + invalidInputError(False, + f"Splitted attention weights should be of size " + f"{attn_weights_split_size}, but is {attn_weights_split.size()}") + + if attention_mask is not None: + attn_mask_size = (bsz, 1, q_len, kv_seq_len) + if attention_mask.size() != attn_mask_size: + invalidInputError(False, + f"Attention mask should be of size {attn_mask_size}, " + f"but is {attention_mask.size()}") + attn_weights_split = attn_weights_split + attention_mask + attn_weights_split = nn.functional.softmax(attn_weights_split, dim=-1) + attn_weights_split = nn.functional.dropout(attn_weights_split, + p=attention_dropout, + training=training) + attn_outputs.append(torch.matmul(attn_weights_split, v)) + attn_output = torch.cat(attn_outputs, dim=1) + return attn_output, None