🔗 Helpful Links

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

• 📄 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

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

✅ No Failures

As of commit 15134ea with merge base 7a0f933 ():
💚 Looks good so far! There are no failures yet. 💚

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

Curious for CP, could also explicit variants / args be added to F.sdpa (e.g. a new backend and a new dict arg backend_hints = {} which could allow passing in any options for CP)? or a new explicit function be introduced? I think in some educational contexts it's better to also have an explicit way of invoking these distributed SDPA directly...

Curious for CP, could also explicit variants / args be added to F.sdpa (e.g. a new backend and a new dict arg backend_hints = {} which could allow passing in any options for CP)? or a new explicit function be introduced? I think in some educational contexts it's better to also have an explicit way of invoking these distributed SDPA directly...

CP currently only supports what F.sdpa and FlexAttention, we basically just forward the arguments to the underlying SDPA or FlexAttention. Also there is a backend selection API for SDPA. Can you describe more your requirement if I misunderstood your request? @vadimkantorov

@pytorchbot merge

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

Can you describe more your requirement if I misunderstood your request?

Well, I was proposing to add an explicit way of invoking CP-SDPA via F.sdpa(..., backend="cp", backend_hints = {...cp_params}) (related in spirit to #137574) or via F.cp_attention(...) (akin to flex_attention) - e.g. so that there is also a way of explicitly invoking CP-SDPA without context manager.