[FA-2 / Mistral] Supprot fa-2 + right padding + forward (huggingf…

…ace#26912) supprot fa-2 + right padding + forward
staghado · Oct 24, 2023 · 8fd2bf3 · 8fd2bf3
1 parent abcfe35
commit 8fd2bf3
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/transformers/models/mistral/modeling_mistral.py b/src/transformers/models/mistral/modeling_mistral.py
@@ -879,6 +879,7 @@ def forward(
             padding_mask is not None
             and hasattr(self.config, "_flash_attn_2_enabled")
             and self.config._flash_attn_2_enabled
+            and past_key_values is not None
         ):
             is_padding_right = padding_mask[:, -1].sum().item() != batch_size
             if is_padding_right: