Skip to content

[MPS] expand the current export unit test for SDPA#177686

Draft
mergennachin wants to merge 1 commit intomainfrom
mergennachin/fix_sdpa_export
Draft

[MPS] expand the current export unit test for SDPA#177686
mergennachin wants to merge 1 commit intomainfrom
mergennachin/fix_sdpa_export

Conversation

@mergennachin
Copy link
Copy Markdown
Contributor

@mergennachin mergennachin commented Mar 17, 2026

Summary

Move the SDPA dynamic seq len export test from test_mps.py to test_aot_inductor.py so it runs on all devices (cpu, gpu, mps), using F.scaled_dot_product_attention instead of the MPS-specific op.

Also clean up the meta registration for _scaled_dot_product_attention_math_for_mps: inline the single-use sdpa_general_mps() helper and remove the dead ensure_4d(key) call, both left over from the branching logic removed in #177620.

Regression test for #177603.

Test plan

  • python test/inductor/test_aot_inductor.py -k test_sdpa_dynamic_seq_len

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @jataylo @hvaara

@mergennachin mergennachin requested a review from malfet as a code owner March 17, 2026 19:39
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot bot commented Mar 17, 2026

🔗 Helpful Links

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

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

❌ 9 New Failures, 1 Unrelated Failure

As of commit e9892cb with merge base 14b40e9 (image):

NEW FAILURES - The following jobs have failed:

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

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

@pytorch-bot pytorch-bot bot added ciflow/mps Run MPS tests (subset of trunk) topic: not user facing topic category labels Mar 17, 2026
Copy link
Copy Markdown
Contributor

@malfet malfet left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's move this whole test to test_aot_inductor.py so that we test it for all devices.
Also, would please consider deleting unused functions from meta registration

Address review feedback: move the SDPA dynamic seq len export test from
test_mps.py to test_aot_inductor.py so it runs on all devices (cpu, gpu,
mps). Use F.scaled_dot_product_attention instead of the MPS-specific op.

Also clean up the meta registration for _scaled_dot_product_attention_math_for_mps:
inline the single-use sdpa_general_mps() helper and remove the dead
ensure_4d(key) call, both left over from the branching logic removed in #177620.

Authored with Claude.
@mergennachin mergennachin force-pushed the mergennachin/fix_sdpa_export branch from 8b29ae7 to e9892cb Compare March 17, 2026 21:54
@mergennachin mergennachin marked this pull request as draft March 18, 2026 11:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants