Update modeling_motif.py

modeling_motif.py

@@ -35,11 +35,7 @@ logger = logging.get_logger(__name__)
 if is_flash_attn_2_available():
     from transformers.modeling_flash_attention_utils import _flash_attention_forward
 
-MorehRMSNorm = None
-ScaledDotProductAttention = None
-MorehFlashAttention = None
 
-#_CHECKPOINT_FOR_DOC = "moreh/Motif-102B"
 _CONFIG_FOR_DOC = "MotifConfig"
 
 from transformers.activations import ACT2CLS as _ACT2CLS
@@ -538,28 +534,27 @@ class MotifFlashAttention2(MotifAttention):
         return tensor.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
 
     def _compute_attention(self, query_states, key_states, value_states, attention_mask, q_len, position_ids,
-                           dropout_rate, sliding_window, is_moreh_attention, batch_num):
+                           dropout_rate, sliding_window, batch_num):
         """Flash Attention 2 implements"""
-        if is_moreh_attention:
-            scale_factor = 1.0 / math.sqrt(self.head_dim)
-            # Copied from _flash_attention_forward
-            if not self._flash_attn_uses_top_left_mask:
-                causal = self.is_causal
-            else:
-                causal = self.is_causal and q_len != 1
+        scale_factor = 1.0 / math.sqrt(self.head_dim)
+        # Copied from _flash_attention_forward
+        if not self._flash_attn_uses_top_left_mask:
+            causal = self.is_causal
+        else:
+            causal = self.is_causal and q_len != 1
 
-            bsz = query_states.shape[0]
+        bsz = query_states.shape[0]
 
-            return _flash_attention_forward(query_states,
-                                            key_states,
-                                            value_states,
-                                            attention_mask,
-                                            q_len,
-                                            position_ids=position_ids,
-                                            dropout=dropout_rate,
-                                            sliding_window=sliding_window,
-                                            is_causal=self.is_causal,
-                                            use_top_left_mask=self._flash_attn_uses_top_left_mask)
+        return _flash_attention_forward(query_states,
+                                        key_states,
+                                        value_states,
+                                        attention_mask,
+                                        q_len,
+                                        position_ids=position_ids,
+                                        dropout=dropout_rate,
+                                        sliding_window=sliding_window,
+                                        is_causal=self.is_causal,
+                                        use_top_left_mask=self._flash_attn_uses_top_left_mask)
 
     def forward(
             self,
@@ -660,13 +655,12 @@ class MotifFlashAttention2(MotifAttention):
         k1, k2 = k1.contiguous(), k2.contiguous()
         v1, v2 = v1.contiguous(), v2.contiguous()
 
-        is_moreh_attention = MorehFlashAttention is not None
-
-        attn11, attn12 = self._compute_attention(q1, k1, v1, attention_mask, q_len, position_ids, dropout_rate, sliding_window, is_moreh_attention, self.batch_num), \
-                            self._compute_attention(q1, k1, v2, attention_mask, q_len, position_ids, dropout_rate, sliding_window, is_moreh_attention, self.batch_num)
-        attn21, attn22 = self._compute_attention(q2, k2, v1, attention_mask, q_len, position_ids, dropout_rate, sliding_window, is_moreh_attention, self.batch_num), \
-                            self._compute_attention(q2, k2, v2, attention_mask, q_len, position_ids, dropout_rate, sliding_window, is_moreh_attention, self.batch_num)
+        attn11, attn12 = self._compute_attention(q1, k1, v1, attention_mask, q_len, position_ids, dropout_rate, sliding_window, self.batch_num), \
+                            self._compute_attention(q1, k1, v2, attention_mask, q_len, position_ids, dropout_rate, sliding_window, self.batch_num)
+        attn21, attn22 = self._compute_attention(q2, k2, v1, attention_mask, q_len, position_ids, dropout_rate, sliding_window, self.batch_num), \
+                            self._compute_attention(q2, k2, v2, attention_mask, q_len, position_ids, dropout_rate, sliding_window, self.batch_num)
 
+        
         attn1, attn2 = torch.cat([attn11, attn12], dim=-1), torch.cat([attn21, attn22], dim=-1)
 
         lambda_q1 = self.lambda_q1.unsqueeze(0).expand([bsz, self.lambda_q1.shape[0]])  # bsz, num_head