Update modeling_motif.py

modeling_motif.py

@@ -571,33 +571,6 @@ class MotifFlashAttention2(MotifAttention):
 
             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:
             return _flash_attention_forward(query_states,
                                             key_states,
                                             value_states,