Ignore non-causal mask in more cases with SDPA #30138
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| @@ -406,31 +406,18 @@ def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, | |||
| tgt_len (`int`): | |||
| The target length or query length the created mask shall have. | |||
| """ | |||
| batch_size, key_value_length = mask.shape | |||
| _, key_value_length = mask.shape | |||
There was a problem hiding this comment.
Can we change the input arg mask: torch.Tensor to mask: Optional[torch.Tensor] and return None immediately if mask is None? The docstring is not compliant with the actual input. (mask ("torch.Tensor" or "None" ):)
Will it break the is_tracing check?
There was a problem hiding this comment.
@minostauros Yes indeed ideally we would want to do that. In practice, the calls to these functions in modeling files are always guarded by:
if attention_mask is not None:but we should IMO indeed accept Optional[torch.Tensor]. I'll leave that to an other PR.
fxmarty
requested review from
amyeroberts,
LysandreJik and
ArthurZucker
and removed request for
ArthurZucker
|
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread. Please note that issues that do not follow the contributing guidelines are likely to be ignored. |
|
This is probably fixed on main |
| is_causal = ( | ||
| True if self.is_decoder and not is_cross_attention and attention_mask is None and tgt_len > 1 else False | ||
| ) |
Fixes #30095. In the non-causal case, we can ignore the mask even for
key_value_length != tgt_lenas we will never have any fully masked row and will never hit the SDPA's mem-efficient attention backend issue.Depends on #28802