Delay and probably avoid unnecessary graph breaks in _upad_input of modeling_flash_attention_utils.py

src/transformers/modeling_flash_attention_utils.py

@@ -205,7 +205,7 @@ def _pad_input(hidden_states, indices, batch, seqlen):
     return output.view(batch, seqlen, *dim)
 
 
-def _get_unpad_data(attention_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]:
+def _get_unpad_data(attention_mask: torch.Tensor) -> tuple[torch.Tensor, int]:
     """
     Retrieves indexing data required to repad unpadded (ragged) tensors.
 
@@ -214,21 +214,15 @@ def _get_unpad_data(attention_mask: torch.Tensor) -> tuple[torch.Tensor, torch.T
             Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid.
 
     Return:
-        indices (`torch.Tensor`):
-            The indices of non-masked tokens from the flattened input sequence.
         cu_seqlens (`torch.Tensor`):
             The cumulative sequence lengths, used to index into ragged (unpadded) tensors. `cu_seqlens` shape is (batch_size + 1,).
         max_seqlen_in_batch (`int`):
             Maximum sequence length in batch.
     """
     seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
-    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
-    # NOTE: Similar to the `.item()` in prepare_fa2_from_position_ids, with torch compile,
-    # this might cause a graph break
-    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    max_seqlen_in_batch = seqlens_in_batch.max()
     cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
     return (
-        indices,
         cu_seqlens,
         max_seqlen_in_batch,
     )
@@ -275,7 +269,8 @@ def _upad_input(
         (max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`tuple[int]`):
             Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence i.e. query, `max_seqlen_in_batch_k` for the source sequence i.e. key/value).
     """
-    indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+    cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+    flatten_mask = attention_mask.view(-1)
 
     # With static caches, the k/v states may be larger than the mask -> we need to slice them to avoid generating garbage
     # It's a bit of an anti-pattern, but otherwise we silently compute wrong attentions scores
@@ -284,13 +279,15 @@ def _upad_input(
 
     batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
 
-    key_layer = _index_first_axis(key_layer, indices_k)
-    value_layer = _index_first_axis(value_layer, indices_k)
+    key_layer = _index_first_axis(key_layer, flatten_mask)
+    value_layer = _index_first_axis(value_layer, flatten_mask)
     if query_length == kv_seq_len:
-        query_layer = _index_first_axis(query_layer, indices_k)
+        query_layer = _index_first_axis(query_layer, flatten_mask)
         cu_seqlens_q = cu_seqlens_k
-        max_seqlen_in_batch_q = max_seqlen_in_batch_k
-        indices_q = indices_k
+        # NOTE: Similar to the `.item()` in prepare_fa2_from_position_ids, with torch compile,
+        # this might cause a graph break
+        max_seqlen_in_batch_q = max_seqlen_in_batch_k.item()
+        indices_q = flatten_mask.nonzero(as_tuple=False).flatten()
     elif query_length == 1:
         max_seqlen_in_batch_q = 1
         cu_seqlens_q = torch.arange(