[inductor] make mix order reduction work with dynamic shapes by shunting314 · Pull Request #168117 · pytorch/pytorch

shunting314 · 2025-11-18T22:20:08Z

Stack from ghstack (oldest at bottom):

[inductor] do mix order reduction fusion earlier and allow fusing more nodes into it #168209
-> [inductor] make mix order reduction work with dynamic shapes #168117

Internal models have dynamic shapes for rms/layer norm. Mix order reduction would be by-passed without this PR.

co-author this PR with Paul Zhang.

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

Differential Revision: D87378475

[ghstack-poisoned]

ghstack-source-id: 3154c25 Pull Request resolved: #168117

pytorch-bot · 2025-11-18T22:20:13Z

🔗 Helpful Links

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

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit 8ab71d2 with merge base 7392106 ():

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

trunk / linux-jammy-py3-clang12-executorch / test (executorch, 1, 1, linux.2xlarge, unstable) (gh) (#166072)
extension/llm/custom_ops/test_sdpa_with_kv_cache.py::SDPAWithAttentionMaskTest::test_sdpa_with_cache_no_mqa_1

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

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben [ghstack-poisoned]

ghstack-source-id: de415ec Pull Request resolved: #168117

shunting314 · 2025-11-18T22:28:15Z

Need a bit more work to properly print dynamic shape expressions

PaulZhang12 · 2025-11-18T22:35:41Z

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

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben Differential Revision: [D87378475](https://our.internmc.facebook.com/intern/diff/D87378475) [ghstack-poisoned]

ghstack-source-id: 62f2d99 Pull Request resolved: #168117

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben Differential Revision: [D87378475](https://our.internmc.facebook.com/intern/diff/D87378475) [ghstack-poisoned]

ghstack-source-id: 98550bc Pull Request resolved: #168117

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben Differential Revision: [D87378475](https://our.internmc.facebook.com/intern/diff/D87378475) [ghstack-poisoned]

ghstack-source-id: 5ac2305 Pull Request resolved: #168117

Internal models have dynamic shapes for rms/layer norm. Mix order reduction would be by-passed without this PR. co-author this PR with Paul Zhang. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben Differential Revision: [D87378475](https://our.internmc.facebook.com/intern/diff/D87378475) [ghstack-poisoned]

ghstack-source-id: df0c686 Pull Request resolved: #168117

PaulZhang12 · 2025-11-19T02:57:48Z

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

PaulZhang12 · 2025-11-19T03:00:24Z

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

PaulZhang12 · 2025-11-19T14:46:20Z

torch/_inductor/codegen/simd.py

+        # since we already add guards in MixOrderReduction.can_fuse
+        # for dynamic shapes.
        if not V.graph.sizevars.statically_known_gt(
            numel,


I don't believe this is true. It can still be the case that this is always false and we have infinite recursion. Better to use the hint here?

Same with the assert below of numel greater than rnumel

Internal models have dynamic shapes for rms/layer norm. Mix order reduction would be by-passed without this PR. co-author this PR with Paul Zhang. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben Differential Revision: [D87378475](https://our.internmc.facebook.com/intern/diff/D87378475) [ghstack-poisoned]

ghstack-source-id: 91c6024 Pull Request resolved: #168117

Internal models have dynamic shapes for rms/layer norm. Mix order reduction would be by-passed without this PR. co-author this PR with Paul Zhang. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben Differential Revision: [D87378475](https://our.internmc.facebook.com/intern/diff/D87378475) [ghstack-poisoned]

ghstack-source-id: 411184b Pull Request resolved: #168117

PaulZhang12 · 2025-11-20T15:36:24Z

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

v0i0 · 2025-11-20T17:14:37Z

torch/_inductor/codegen/simd.py

-            numel,
-            rnumel,
-        )
+        assert V.graph.sizevars.evaluate_expr(sympy.Gt(numel, rnumel))


just for my information as someone who's new: How are those expressions different? i guess one doesn't force guarding?

right. For statically_known_gt, we return false if we can not decide due to dynamic shapes. For evaluate_expr, we will add guards and recompile if the guard fails.

shunting314 · 2025-11-20T21:52:45Z

@pytorchbot merge

pytorchmergebot · 2025-11-20T21:54:36Z

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

eellison · 2025-11-20T23:58:42Z

torch/_inductor/codegen/simd.py

-            numel,
-            rnumel,
-        ):
+        if not V.graph.sizevars.evaluate_expr(sympy.Gt(numel, rnumel)):


do we not want to guard here ?

I'm a bit conservative to adding guards. In real use cases we see, numel may have dynamic shape (e.g. jagged form of inputs), but rnumel is always static shape decided by model parameters

eellison · 2025-11-20T23:59:02Z

torch/_inductor/choices.py

-        if (
-            shared_data_score < config.score_fusion_memory_threshold
-        ) and not MixOrderReduction.can_fuse(node1, node2):
+        if MixOrderReduction.can_fuse(node1, node2):


for my knowledge - why do we not care about shared dat a?

score_fusion_memory returns 0 for MixOrderReduction (the next PR fixes that) so this PR skips the check. Not a big deal since we know mix order reduction always have decent shared data. (otherwise we don't fuse)

eellison · 2025-11-20T23:59:33Z

torch/_inductor/scheduler.py

+        # Call evaluate_expr rather than statically_known_geq since nrow can
+        # have dynamic shape in real models.
+        # Don't use hint directly since hint can be non-representative.
+        if not V.graph.sizevars.evaluate_expr(sympy.Ge(nrow * ncol, size_thres)):


similarly - wondering why not to guard?

This one adds guard. nrow may contains dynamic shapes.

Internal models have dynamic shapes for rms/layer norm. Mix order reduction would be by-passed without this PR. co-author this PR with Paul Zhang. Differential Revision: [D87378475](https://our.internmc.facebook.com/intern/diff/D87378475) Pull Request resolved: #168117 Approved by: https://github.com/v0i0, https://github.com/jansel

[inductor] make mix order reduction work with dynamic shapes

081144b

[ghstack-poisoned]

shunting314 added a commit that referenced this pull request Nov 18, 2025

[inductor] make mix order reduction work with dynamic shapes

a23199f

ghstack-source-id: 3154c25 Pull Request resolved: #168117

pytorch-bot bot added ciflow/inductor module: inductor labels Nov 18, 2025

Update on "[inductor] make mix order reduction work with dynamic shapes"

bf0e0ce

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben [ghstack-poisoned]

shunting314 added a commit that referenced this pull request Nov 18, 2025

[inductor] make mix order reduction work with dynamic shapes

7af92a7

ghstack-source-id: de415ec Pull Request resolved: #168117

shunting314 marked this pull request as draft November 18, 2025 22:27

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

shunting314 added a commit that referenced this pull request Nov 18, 2025

[inductor] make mix order reduction work with dynamic shapes

5c20269

ghstack-source-id: 62f2d99 Pull Request resolved: #168117

shunting314 added a commit that referenced this pull request Nov 19, 2025

[inductor] make mix order reduction work with dynamic shapes

8be25be

ghstack-source-id: 98550bc Pull Request resolved: #168117

shunting314 added a commit that referenced this pull request Nov 19, 2025

[inductor] make mix order reduction work with dynamic shapes

f824e23

ghstack-source-id: 5ac2305 Pull Request resolved: #168117

shunting314 marked this pull request as ready for review November 19, 2025 00:45

shunting314 added a commit that referenced this pull request Nov 19, 2025

[inductor] make mix order reduction work with dynamic shapes

008ce73

ghstack-source-id: df0c686 Pull Request resolved: #168117

PaulZhang12 reviewed Nov 19, 2025

View reviewed changes

shunting314 added a commit that referenced this pull request Nov 19, 2025

[inductor] make mix order reduction work with dynamic shapes

c28a534

ghstack-source-id: 91c6024 Pull Request resolved: #168117

shunting314 added a commit that referenced this pull request Nov 19, 2025

[inductor] make mix order reduction work with dynamic shapes

9bee762

ghstack-source-id: 411184b Pull Request resolved: #168117

shunting314 mentioned this pull request Nov 20, 2025

[inductor] do mix order reduction fusion earlier and allow fusing more nodes into it #168209

Closed

shunting314 requested review from eellison and jansel November 20, 2025 01:17

shunting314 requested a review from v0i0 November 20, 2025 01:17

v0i0 reviewed Nov 20, 2025

View reviewed changes

v0i0 approved these changes Nov 20, 2025

View reviewed changes

jansel approved these changes Nov 20, 2025

View reviewed changes

shunting314 added the topic: not user facing topic category label Nov 20, 2025

pytorchmergebot added the merging label Nov 20, 2025

pytorchmergebot closed this in 6644fd7 Nov 20, 2025

pytorchmergebot added Merged and removed merging labels Nov 20, 2025

eellison reviewed Nov 20, 2025

View reviewed changes

shunting314 mentioned this pull request Nov 21, 2025

[inductor] fix picking wrong contiguous node #168371

Closed

github-actions bot deleted the gh/shunting314/266/head branch December 21, 2025 02:20

Conversation

shunting314 commented Nov 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Nov 18, 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/168117

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

shunting314 commented Nov 18, 2025

Uh oh!

PaulZhang12 commented Nov 18, 2025

Uh oh!

PaulZhang12 commented Nov 19, 2025

Uh oh!

PaulZhang12 commented Nov 19, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

PaulZhang12 commented Nov 20, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shunting314 commented Nov 20, 2025

Uh oh!

pytorchmergebot commented Nov 20, 2025

Merge started

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shunting314 Nov 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

shunting314 commented Nov 18, 2025 •

edited

Loading

pytorch-bot bot commented Nov 18, 2025 •

edited

Loading

shunting314 Nov 21, 2025 •

edited

Loading