Skip to content

Commit

Permalink
prepare_fa2_from_position_ids function bugfix (#33269)
Browse files Browse the repository at this point in the history
contiguous() is called before view() for key and value within prepare_fa2_from_position_ids function
  • Loading branch information
meliksahturker authored Nov 25, 2024
1 parent a0f4f31 commit c50b567
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions src/transformers/modeling_flash_attention_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -163,8 +163,8 @@ def prepare_fa2_from_position_ids(query, key, value, position_ids):
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).
"""
query = query.view(-1, query.size(-2), query.size(-1))
key = key.view(-1, key.size(-2), key.size(-1))
value = value.view(-1, value.size(-2), value.size(-1))
key = key.contiguous().view(-1, key.size(-2), key.size(-1))
value = value.contiguous().view(-1, value.size(-2), value.size(-1))
position_ids = position_ids.flatten()
indices_q = torch.arange(position_ids.size(0), device=position_ids.device, dtype=torch.int32)

Expand Down

0 comments on commit c50b567

Please sign in to comment.