Important: Update modeling_mpt.py after Flash Attention 2.7.0

modeling_mpt.py

@@ -140,9 +140,9 @@ def gen_flash_attn_padding_info(bsz: int, S: int, past_key_len: int, device: tor
         key_padding_mask = attention_mask_in_length
         query_padding_mask = attention_mask_in_length
         unpadding_function = bert_padding.unpad_input_for_concatenated_sequences
-    (_, indices_q, cu_seqlens_q, max_seqlen_q) = unpadding_function(torch.empty(bsz, S, 1, device=device), query_padding_mask)
-    (_, indices_k, cu_seqlens_k, max_seqlen_k) = unpadding_function(torch.empty(bsz, past_key_len + S, 1, device=device), key_padding_mask)
-    (_, indices_v, _, _) = unpadding_function(torch.empty(bsz, past_key_len + S, 1, device=device), key_padding_mask)
+    (_, indices_q, cu_seqlens_q, max_seqlen_q, *rest) = unpadding_function(torch.empty(bsz, S, 1, device=device), query_padding_mask)
+    (_, indices_k, cu_seqlens_k, max_seqlen_k, *rest) = unpadding_function(torch.empty(bsz, past_key_len + S, 1, device=device), key_padding_mask)
+    (_, indices_v, _, _, *rest) = unpadding_function(torch.empty(bsz, past_key_len + S, 1, device=device), key_padding_mask)
     flash_attn_padding_info['indices_q'] = indices_q
     flash_attn_padding_info['indices_k'] = indices_k
     flash_attn_padding_info['indices_v'] = indices_v