[inductor][triton pin] add support for new TMA API for mm.py templates by davidberard98 · Pull Request #155723 · pytorch/pytorch

davidberard98 · 2025-06-11T17:54:14Z

Stack from ghstack (oldest at bottom):

Triton 3.4 will remove the experimental TMA APIs: triton-lang/triton#6488

For mm.py templates, this PR adds support for using the new APIs when they are available (and otherwise falls back to the experimental APIs).

For flex_attention, we'll remove TMA support for Triton 3.2 and 3.3 (versions of triton that don't have the new API).

For mm_scaled_grouped.py, #150944 will remove TMA support for Triton 3.2.

Note: we attempted this earlier with #154858, but this broke TMA usage in Triton 3.2.

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

Differential Revision: D76444471

…emplates [ghstack-poisoned]

pytorch-bot · 2025-06-11T17:54:18Z

🔗 Helpful Links

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

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

✅ You can merge normally! (2 Unrelated Failures)

As of commit cd4f19f with merge base 717a099 ():

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

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

pull / cuda12.8-py3.10-gcc9-sm75 / test (pr_time_benchmarks, 1, 1, linux.g4dn.metal.nvidia.gpu) (gh) (trunk failure)
MISSING REGRESSION TEST
pull / linux-jammy-py3-clang12-executorch / test (executorch, 1, 1, linux.2xlarge) (gh) (trunk failure)
exir/backend/test/test_to_backend_multi_method.py::TestToBackendMultiMethod::test_multi_method_end_to_end

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

davidberard98 · 2025-06-11T18:13:06Z

@davidberard98 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

NikhilAPatel · 2025-06-11T21:36:43Z

I'm not entirely familiar with the mechanics of AOTI, so this may be wrong. But since AOTI doesn't support the new TMA API (due to trouble setting the allocator function), would this cause failures if our Triton version is high enough such that has_triton_stable_tma_api() returns True but we're using AOTI? Because then we'd choose the new API codepath but we wouldn't have an allocator function set.

Maybe for now, has_triton_stable_tma_api() should also make sure we're not using AOTI?

…py templates" Triton 3.4 will remove the experimental TMA APIs: triton-lang/triton#6488 For mm.py templates, this PR adds support for using the new APIs when they are available (and otherwise falls back to the experimental APIs). For flex_attention, we'll remove TMA support for Triton 3.2 and 3.3 (versions of triton that don't have the new API). For mm_scaled_grouped.py, #150944 will remove TMA support for Triton 3.2. Note: we attempted this earlier with #154858, but this broke TMA usage in Triton 3.2. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov Differential Revision: [D76444471](https://our.internmc.facebook.com/intern/diff/D76444471) [ghstack-poisoned]

…emplates ghstack-source-id: 543fae4 Pull Request resolved: #155723

davidberard98 · 2025-06-11T22:59:36Z

@NikhilAPatel good catch - I added a check for that (and ran some really quick sanity tests to make sure that AOTI works w/ new triton, using python test/inductor/test_aot_inductor.py -k mm)

davidberard98 · 2025-06-11T23:00:28Z

@davidberard98 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

NikhilAPatel · 2025-06-11T23:56:57Z

torch/_inductor/utils.py

        inner_bytes = inner_dim * dtype.itemsize
        return V.graph.sizevars.statically_known_multiple_of(inner_bytes, TMA_ALIGNMENT)

+    if has_triton_stable_tma_api() and config.cpp_wrapper:


Will this disable TMA for AOTI always? Because we still want TMA for AOTI as long as its using the old API right?

has_triton_stable_tma_api() means "we'll be using the new API"

Oh true my bad

davidberard98 · 2025-06-12T00:11:03Z

@pytorchbot merge

pytorchmergebot · 2025-06-12T00:13:05Z

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 · 2025-06-12T00:13:16Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

pull / cuda12.8-py3.10-gcc9-sm75 / build

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

davidberard98 · 2025-06-12T00:50:19Z

@pytorchbot merge

pytorchmergebot · 2025-06-12T00:52:15Z

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 · 2025-06-12T00:52:26Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

pull / cuda12.8-py3.10-gcc9-sm75 / build

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

davidberard98 · 2025-06-12T01:37:46Z

@pytorchbot merge

pytorchmergebot · 2025-06-12T01:39:44Z

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

Summary: Add utils function which verifies that `triton.tools.experimental_descriptor` is available in the current Triton version. This diff expands the work in this [GitHub PR](pytorch/pytorch#155723), which finds that Triton 3.4.0+ deprecated the experimental descriptor. Differential Revision: D82551608

[inductor][triton pin] use experimental or stable TMA API for mm.py t…

e4de734

…emplates [ghstack-poisoned]

pytorch-bot bot added ciflow/inductor module: inductor labels Jun 11, 2025

davidberard98 changed the title ~~[inductor][triton pin] use experimental or stable TMA API for mm.py templates~~ [inductor][triton pin] add support for new TMA API for mm.py templates Jun 11, 2025

davidberard98 added the release notes: inductor label Jun 11, 2025

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jun 11, 2025

davidberard98 added the ciflow/h100 label Jun 11, 2025

davidberard98 requested review from NikhilAPatel, PaulZhang12 and aakhundov June 11, 2025 21:25

davidberard98 added a commit that referenced this pull request Jun 11, 2025

[inductor][triton pin] use experimental or stable TMA API for mm.py t…

5819f31

…emplates ghstack-source-id: 543fae4 Pull Request resolved: #155723

NikhilAPatel reviewed Jun 11, 2025

View reviewed changes

NikhilAPatel approved these changes Jun 12, 2025

View reviewed changes

pytorchmergebot added the merging label Jun 12, 2025

pytorchmergebot removed the merging label Jun 12, 2025

pytorchmergebot added the merging label Jun 12, 2025

pytorchmergebot removed the merging label Jun 12, 2025

pytorchmergebot added the merging label Jun 12, 2025

davidberard98 mentioned this pull request Jun 12, 2025

[flex attention][triton pin] use new TMA API #155771

Closed

pytorchmergebot added the Merged label Jun 12, 2025

pytorchmergebot closed this in c3ecabf Jun 12, 2025

pytorchmergebot removed the merging label Jun 12, 2025

davidberard98 mentioned this pull request Jun 27, 2025

Triton has removed the experimental descriptor API #154162

Closed

github-actions bot deleted the gh/davidberard98/372/head branch July 14, 2025 02:20

jananisriram mentioned this pull request Sep 16, 2025

Verify experimental_descriptor in Triton version meta-pytorch/tritonbench#428

Merged

Conversation

davidberard98 commented Jun 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jun 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ You can merge normally! (2 Unrelated Failures)

Uh oh!

davidberard98 commented Jun 11, 2025

Uh oh!

NikhilAPatel commented Jun 11, 2025

Uh oh!

davidberard98 commented Jun 11, 2025

Uh oh!

davidberard98 commented Jun 11, 2025

Uh oh!

NikhilAPatel Jun 11, 2025

Choose a reason for hiding this comment

Uh oh!

davidberard98 Jun 12, 2025

Choose a reason for hiding this comment

Uh oh!

NikhilAPatel Jun 12, 2025

Choose a reason for hiding this comment

Uh oh!

davidberard98 commented Jun 12, 2025

Uh oh!

pytorchmergebot commented Jun 12, 2025

Merge started

Uh oh!

pytorchmergebot commented Jun 12, 2025

Merge failed

Uh oh!

davidberard98 commented Jun 12, 2025

Uh oh!

pytorchmergebot commented Jun 12, 2025

Merge started

Uh oh!

pytorchmergebot commented Jun 12, 2025

Merge failed

Uh oh!

davidberard98 commented Jun 12, 2025

Uh oh!

pytorchmergebot commented Jun 12, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

davidberard98 commented Jun 11, 2025 •

edited

Loading

pytorch-bot bot commented Jun 11, 2025 •

edited

Loading