Causal masking hint

marc_stovell: <p>I am surprised by this piece of code in <a href="https://github.com/pytorch/pytorch/blob/main/torch/nn/functional.py#L6252" rel="noopener nofollow ugc">functional.py</a></p>
<pre><code class="lang-auto">    if is_causal and attn_mask is None:
        raise RuntimeError(
            "Need attn_mask if specifying the is_causal hint. "
            "You may use the Transformer module method "
            "`generate_square_subsequent_mask` to create this mask."
        )

    if is_causal and key_padding_mask is None and not need_weights:
        # when we have a kpm or need weights, we need attn_mask
        # Otherwise, we use the is_causal hint go as is_causal
        # indicator to SDPA.
        attn_mask = None
</code></pre>
<p>The first test requires an attention mask when is_causal == True. The second test says that we do not need it in some subcase (is_causal and key_padding_mask is None and not need_weights). So if you are in the subcase, you have to pass the attention mask but it will be set to None <img src="https://discuss.pytorch.org/images/emoji/apple/exploding_head.png?v=14" title=":exploding_head:" class="emoji" alt=":exploding_head:" loading="lazy" width="20" height="20"></p>
<p>I was expecting to not provide an attention mask if I want pytorch to create a causal mask for me.</p>
<p>Is it a bug or I am missing something?</p>
<p>Thank you!</p>

vdw: <p>According to the code comment:</p>
<pre><code class="lang-auto"> need_weights: output attn_output_weights.
            Default: `True`
            Note: `needs_weight` defaults to `True`, but should be set to `False`
            For best performance when attention weights are not needed.
            *Setting needs_weights to `True`
            leads to a significant performance degradation.*
</code></pre>
<p>Masking is done to ignore certain elements in the attention matrix. If the attention weights are not used at all, no need for masking.</p>
<p>That being said, I can’t tell case where you don’t need the weights.</p>

Arunprakash-A: <p>That’s a good observation.</p>
<p>PyTorch recommends setting <code>need_weights=False</code> (the default is <em>True</em>) for improved performance<a href="https://github.com/pytorch/pytorch/blob/main/torch/nn/functional.py#L6126" rel="noopener nofollow ugc">[source]</a>. When setting this to <code>need_weights=False</code> and <em><code>Keep_padding_mask is None</code></em> , the subcase is executed</p>
<p>In this case, as noted in the comments, PyTorch uses a more <a href="https://github.com/pytorch/pytorch/blob/main/torch/nn/functional.py#L5805" rel="noopener nofollow ugc">efficient implementation</a> of scaled dot product attention (SDPA) by fusing multiple operations.</p>
<p>This efficient implementation requires <code>attention_mask=None</code> (even though masking is still handled internally). Therefore, even if one passes the <code>attention_mask</code> manually, it gets overridden.</p>
<p>Why not always use this fused version, then?<br>
Likely because fusion prevents returning the attention weights. For more details, refer to <a href="https://github.com/Dao-AILab/flash-attention" rel="noopener nofollow ugc">FlashAttention.</a></p>
<p>Hope it helps!</p>