fix flash_attn input

PaddlePaddle · Feb 24, 2025 · 03d6a0b · 03d6a0b
1 parent fd59b6b
commit 03d6a0b
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/paddlenlp/experimental/transformers/fused_transformer_layers.py b/paddlenlp/experimental/transformers/fused_transformer_layers.py
@@ -2987,8 +2987,8 @@ def compute_mla_absorb(
                 value,
                 kwargs.get("cu_seqlens_q", None),
                 kwargs.get("cu_seqlens_k", None),
-                kwargs.get("max_input_length", -1),
-                kwargs.get("max_input_length", -1),
+                kwargs.get("max_enc_len_this_time", -1),
+                kwargs.get("max_enc_len_this_time", -1),
                 self.softmax_scale,
                 causal=True,
                 training=False,
@@ -3316,8 +3316,8 @@ def compute_mla_absorb(
                 value,
                 kwargs.get("cu_seqlens_q", None),
                 kwargs.get("cu_seqlens_k", None),
-                kwargs.get("max_input_length", -1),
-                kwargs.get("max_input_length", -1),
+                kwargs.get("max_enc_len_this_time", -1),
+                kwargs.get("max_enc_len_this_time", -1),
                 self.softmax_scale,
                 causal=True,
                 training=False,