remove causal mask

006e8c6a · Atream · cdb6f896 · 006e8c6a
Commit 006e8c6a authored Feb 23, 2025 by Atream
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 3 deletions

ktransformers/operators/models.py ktransformers/operators/models.py +6 -3

No files found.
--- a/ktransformers/operators/models.py
+++ b/ktransformers/operators/models.py
@@ -649,9 +649,12 @@ class KDeepseekV2Model(BaseInjectedModule):
        if per_layer_prefill_flag:
            causal_mask = None
        else:
-            causal_mask = self._update_causal_mask(
-                attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
-            )
+            if os.name == 'nt':
+                causal_mask = self._update_causal_mask(
+                    attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+                )
+            else:
+                causal_mask = None

        # embed positions
        hidden_states = inputs_embeds