LLM export pass to swap in custom SDPA #10355
Conversation
sxu
requested review from
iseeyuan,
jackzhxng,
larryliu0820 and
swolchok
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/10355
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 40181b2 with merge base 22ba09e ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
|
This pull request was exported from Phabricator. Differential Revision: D73444078 |
Summary: Pull Request resolved: pytorch#10355 Differential Revision: D73444078
|
This pull request was exported from Phabricator. Differential Revision: D73444078 |
Summary: Pull Request resolved: pytorch#10355 Differential Revision: D73444078
|
This pull request was exported from Phabricator. Differential Revision: D73444078 |
Summary: Pull Request resolved: pytorch#10355 Differential Revision: D73444078
|
This pull request was exported from Phabricator. Differential Revision: D73444078 |
| def call_operator(self, op, args, kwargs, meta): | ||
| from executorch.extension.llm.custom_ops import custom_ops # noqa | ||
|
|
||
| if op != torch.ops.aten.scaled_dot_product_attention.default: |
There was a problem hiding this comment.
Isn't this op getting decomposed?
There was a problem hiding this comment.
The idea is to run this pass before to_edge, and avoid the decomposed version for perf reasons.
| kT = self._transpose(k, meta) | ||
| vT = self._transpose(v, meta) | ||
|
|
||
| if mask is not None and mask.node.meta["val"].dtype == torch.bool: |
There was a problem hiding this comment.
Put a todo here that custom sdpa once supports boolean mask, this wont be needed. tag me on the todo
| (mask, 0.0, float("-inf")), | ||
| {}, | ||
| meta, | ||
| ) |
There was a problem hiding this comment.
Also worth checking if the mask is > 2D than add appropriate squeeze ops while making sure first N - 2 dims are all 1
| meta, | ||
| ) | ||
|
|
||
| custom_sdpa = super().call_operator( |
There was a problem hiding this comment.
I would like to add option here that allows us to assume that the mask will be causal and thus we can just set mask =None and is_causal = True, can you do that and add corresponding test?
kimishpatel
left a comment
kimishpatel
left a comment
There was a problem hiding this comment.
Some nits and special handling for mask and is_causal requested
Summary: Pull Request resolved: pytorch#10355 Reviewed By: billmguo Differential Revision: D73444078
|
This pull request was exported from Phabricator. Differential Revision: D73444078 |
|
@kimishpatel can you take another look? |
kimishpatel
left a comment
kimishpatel
left a comment
There was a problem hiding this comment.
Thanks for the changes. Looks good. @guangy10 we should look into adopting this as well.
Summary: Pull Request resolved: pytorch#10355 Reviewed By: billmguo, kimishpatel Differential Revision: D73444078
|
This pull request was exported from Phabricator. Differential Revision: D73444078 |
4 participants
Differential Revision: D73444078