Update modeling_motif.py

remove is_moreh_attention, batch_num

modeling_motif.py

@@ -396,11 +396,7 @@ class MotifAttention(nn.Module):
         self.rope_theta = config.rope_theta
         self.is_causal = True
         self.attention_dropout = config.attention_dropout
-        try:
-            self.batch_num = config.batch_num
-            logger.info(f'self.batcn_num : {self.batch_num}')
-        except:
-            self.batch_num = None
+    
         if (self.head_dim * self.num_heads) != self.hidden_size:
             raise ValueError(f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
                              f" and `num_heads`: {self.num_heads}).")
@@ -556,7 +552,7 @@ 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):
         """Flash Attention 2 implements"""
         
         scale_factor = 1.0 / math.sqrt(self.head_dim)
@@ -566,37 +562,7 @@ class MotifFlashAttention2(MotifAttention):
         else:
             causal = self.is_causal and q_len != 1
     
-        if is_moreh_attention:
-            bsz = query_states.shape[0]
-
-            if batch_num:
-                query_states = query_states.reshape(bsz*q_len,self.num_heads,self.head_dim)
-                key_states = key_states.reshape(bsz*q_len,self.num_heads,self.head_dim)
-                value_states = value_states.reshape(bsz*q_len,self.num_heads,self.head_dim)
-
-                attn_out = moreh_ops.flash_attention_varlen_dp(query_states,
-                                                        key_states,
-                                                        value_states,
-                                                        attention_mask,
-                                                        attention_mask,
-                                                        max_seqlen_q=q_len,
-                                                        max_seqlen_kv=q_len,
-                                                        dropout_p=dropout_rate,
-                                                        softmax_scale=scale_factor,
-                                                        is_causal=causal,
-                                                        batch_num=batch_num)
-                attn_out = attn_out.reshape(bsz, q_len, self.num_heads, -1)
-            else:
-                return MorehFlashAttention(query_states,
-                                        key_states,
-                                        value_states,
-                                        padding_mask=attention_mask,
-                                        dropout_p=dropout_rate,
-                                        softmax_scale=scale_factor,
-                                        causal=causal)
-            return attn_out
-        else:
-            attn_out = _flash_attention_forward(query_states.bfloat16(),
+        attn_out = _flash_attention_forward(query_states.bfloat16(),
                                             key_states.bfloat16(),
                                             value_states.bfloat16(),
                                             attention_mask,
@@ -607,8 +573,7 @@ class MotifFlashAttention2(MotifAttention):
                                             is_causal=True,
                                             softmax_scale=scale_factor,
                                             use_top_left_mask=self._flash_attn_uses_top_left_mask)
-            #logger.info(attn_out)
-            return attn_out.float()
+        return attn_out.float()
 
     def forward(
             self,
@@ -709,12 +674,10 @@ 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._compute_attention(q1, k1, v2, attention_mask, q_len, position_ids, dropout_rate, sliding_window)
+        attn21, attn22 = self._compute_attention(q2, k2, v1, attention_mask, q_len, position_ids, dropout_rate, sliding_window), \
+                            self._compute_attention(q2, k2, v2, attention_mask, q_len, position_ids, dropout_rate, sliding_window)
 
         attn1, attn2 = torch.cat([attn11, attn12], dim=-1), torch.cat([attn21, attn22], dim=-1)