[SDPA] update type hint for scaled_dot_product_attention and documentation

[drisspg]

Summary

• Adds type hinting support for SDPA
• Updates the documentation adding warnings and notes on the context manager
• Adds scaled_dot_product_attention to the non-linear activation function section of nn.functional docs

cc @svekars @carljparker

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/94008

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❗ 2 Active SEVs

There are 2 currently active SEVs. If your PR is affected, please view them below:

• ROCm jobs fail to access AMD apt repo
• MacOS M1 CPU/MPS tests are backlogged due to HW availability

✅ No Failures

As of commit 2a93079:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[pytorch-bot]

This PR needs to be approved by an authorized maintainer before merge.

[drisspg]

\text{attn\_mask} = $torch.ones(L,S, torch.bool, diagonal=0).tril()$ if is\_causal

\text{attn\_mask} = $torch.masked\_fill(!attn\_mask, -\infty)$ if attn\_mask.dtype==torch.bool

\text{attn\_weight} = $\text{torch.softmax}(\frac{QK^T}{\sqrt{d^k}}+attn\_mask)$ 

\text{attn\_weight} = $torch.dropout(\text{attn\_weight}, dropout_p)$

\text{return} $torch.matmul(\text{attn\_weight},V)$

Curious if we think this latex math should be added to the longish docstring to explain the math fallback

[drisspg]

Should I add a warn once to scaled_dot_product_efficient_attention, and scaled_dot_product_flash_attention saying:
" Memory-efficeint/FlashAttention SDPA is a beta feature. See the documentation for torch.nn.scaled_dot_product_attention for further information"

warn onces are kind of annoying but we might want to very directly point uses to the docs

[jbschlosser]

Nice work! I added the usual nitpicky grammar / formatting comments for public-facing docs, nothing major.

FYI there's a style guide here for consistency in formatting module docs. I realize this is not a module, but maybe some of the content there will be useful for maintaining consistency.

[drisspg]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

The merge job was canceled. If you believe this is a mistake,then you can re trigger it through pytorch-bot.

[drisspg]

@pytorchbot merge -f "all checks are passing"

[pytorchmergebot]

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here