Add flash attention

modeling_baichuan.py

@@ -138,20 +138,36 @@ class BaichuanAttention(torch.nn.Module):
 
         past_key_value = (key_states, value_states) if use_cache else None
 
-        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
-
-        if attention_mask is not None:
-            if q_len == 1: # inference with cache
-                if len(attention_mask.size()) == 4:
-                    attention_mask = attention_mask[:, :, -1:, :]   
-                else:
-                    attention_mask = attention_mask[:, -1:, :]    
-            attn_weights = attn_weights + attention_mask
-            attn_weights = torch.max(attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min))
+        pytorch_major_version = int(torch.__version__.split('.')[0])
+        if pytorch_major_version >= 2:
+            if attention_mask is not None:
+                if q_len == 1: # inference with cache
+                    if len(attention_mask.size()) == 4:
+                        attention_mask = attention_mask[:, :, -1:, :]   
+                    else:
+                        attention_mask = attention_mask[:, -1:, :]
 
-        attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)
+            attn_output = torch.nn.functional.scaled_dot_product_attention(query_states,
+                                                                           key_states,
+                                                                           value_states,
+                                                                           dropout_p=0.0,
+                                                                           attn_mask=attention_mask)
+        else:
 
-        attn_output = torch.matmul(attn_weights, value_states)
+            attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+    
+            if attention_mask is not None:
+                if q_len == 1: # inference with cache
+                    if len(attention_mask.size()) == 4:
+                        attention_mask = attention_mask[:, :, -1:, :]   
+                    else:
+                        attention_mask = attention_mask[:, -1:, :]    
+                attn_weights = attn_weights + attention_mask
+                attn_weights = torch.max(attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min))
+    
+            attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)
+    
+            attn_output = torch.matmul(attn_weights, value_states)
 
         attn_output = attn_output.transpose(1, 2)
         attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)