Motif-Technologies
/

Motif-2.6B

Text Generation

text-generation-inference

Model card Files Files and versions

leejunhyeok commited on 22 days ago

Commit

20b97f1

·

verified ·

1 Parent(s): 8855d03

Update modeling_motif.py

Files changed (1) hide show

modeling_motif.py +13 -10

modeling_motif.py CHANGED Viewed

@@ -545,16 +545,19 @@ class MotifFlashAttention2(MotifAttention):
         bsz = query_states.shape[0]
-        return _flash_attention_forward(query_states.bfloat16(),
-                                        key_states.bfloat16(),
-                                        value_states.bfloat16(),
-                                        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,

         bsz = query_states.shape[0]
+        return map(
+            lambda x: x.float32(),
+            _flash_attention_forward(query_states.bfloat16(),
+                                key_states.bfloat16(),
+                                value_states.bfloat16(),
+                                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,