[graph partition] Add way to register custom rule #163310

zou3519 · 2025-09-19T01:56:34Z

Stack from ghstack (oldest at bottom):

This PR adds an experimental way to register a custom rule for if
inductor should partition the graph around an operator.

Test Plan:

new test

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

This PR adds an experimental way to register a custom rule for if inductor should partition the graph around an operator. Test Plan: - new test [ghstack-poisoned]

pytorch-bot · 2025-09-19T01:56:38Z

🔗 Helpful Links

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

📄 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, 1 Pending

As of commit d13ce51 with merge base f4c33cd ():
💚 Looks good so far! There are no failures yet. 💚

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

This PR adds an experimental way to register a custom rule for if inductor should partition the graph around an operator. Test Plan: - new test ghstack-source-id: 6627a21 Pull Request resolved: #163310

This PR adds an experimental way to register a custom rule for if inductor should partition the graph around an operator. Test Plan: - new test cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben [ghstack-poisoned]

This PR adds an experimental way to register a custom rule for if inductor should partition the graph around an operator. Test Plan: - new test ghstack-source-id: e1ac4aa Pull Request resolved: #163310

ProExpertProg

This looks good for dynamically deciding whether to partition nodes for custom ops! The only other thing here would be data-dependent dynamic sizes

eellison · 2025-09-19T18:13:09Z

torch/_inductor/scheduler.py

+    `register_should_partition_rule` is currently private and experimental.
+    Use at your own risk.
+    """
+    assert isinstance(op, (torch._ops.OpOverload, torch._ops.HigherOrderOperator))


When would this be torch._ops.HigherOrderOperator ?

no clue lol. It is very likely that the HOP part of this doesn't work

I am going to assert against HOP to start. There's not many HOPs that would be here anyways (they're not FallbackKernel)

eellison · 2025-09-19T18:15:57Z

torch/_inductor/scheduler.py

+                should_partition_fn = _custom_should_partition_fns[operator]
+                fx_node = ir_node.get_origin_node()
+                assert fx_node is not None
+                _, fake_args, fake_kwargs = (


nit, this wont always succeed.. assert on first returned arg ? tbh not sure why i didnt just make this optional return..

Oh I didn't realize the first returned arg was a bool for if this succeeded or not

eellison · 2025-09-19T18:20:03Z

torch/_inductor/scheduler.py

+                should_partition_fn = _custom_should_partition_fns[operator]
+                fx_node = ir_node.get_origin_node()
+                assert fx_node is not None
+                _, fake_args, fake_kwargs = (


Nit: if the operator doesnt have fixed require strides, it's possible the strides of the inputs that will be passed in at runtime are different than the strides in fake_args fake_kwargs. Not sure it really matters though, non-blocking for land.

@IvanKobzarev has an example showing how to get the args with correct strides here: https://github.com/pytorch/pytorch/pull/162377/files#diff-8dfc175f817a04eb10c183509564129a91f32db31190e7b4ad3fd9653337c6f6R3672. but again not landing blocking.

eellison · 2025-09-19T18:20:37Z

torch/_inductor/scheduler.py

 _P = ParamSpec("_P")


+_custom_should_partition_fns: weakref.WeakKeyDictionary[


should there be anyway of resetting this ? it's fine if not since this is very particular use case and vllm can do it manually. also note this wont be part of fx graph cache but i guess that's fine.

easiest way is torch._inductor.scheduler._custom_should_partition_fns.reset() lol

I'm not sure

This PR adds an experimental way to register a custom rule for if inductor should partition the graph around an operator. Test Plan: - new test cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov coconutruben [ghstack-poisoned]

This PR adds an experimental way to register a custom rule for if inductor should partition the graph around an operator. Test Plan: - new test ghstack-source-id: 98eb0f1 Pull Request resolved: #163310

zou3519 · 2025-09-19T23:20:34Z

@pytorchbot merge

pytorchmergebot · 2025-09-19T23:22:22Z

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

This PR adds an experimental way to register a custom rule for if inductor should partition the graph around an operator. Test Plan: - new test Pull Request resolved: #163310 Approved by: https://github.com/ProExpertProg, https://github.com/BoyuanFeng, https://github.com/eellison ghstack dependencies: #162117, #162307, #162651

This PR adds an experimental way to register a custom rule for if inductor should partition the graph around an operator. Test Plan: - new test Pull Request resolved: pytorch#163310 Approved by: https://github.com/ProExpertProg, https://github.com/BoyuanFeng, https://github.com/eellison ghstack dependencies: pytorch#162117, pytorch#162307, pytorch#162651

This PR adds an experimental way to register a custom rule for if inductor should partition the graph around an operator. Test Plan: - new test Pull Request resolved: #163310 Approved by: https://github.com/ProExpertProg, https://github.com/BoyuanFeng, https://github.com/eellison ghstack dependencies: #162117, #162307, #162651

This PR adds an experimental way to register a custom rule for if inductor should partition the graph around an operator. Test Plan: - new test Pull Request resolved: pytorch#163310 Approved by: https://github.com/ProExpertProg, https://github.com/BoyuanFeng, https://github.com/eellison ghstack dependencies: pytorch#162117, pytorch#162307, pytorch#162651

[graph partition] Add way to register custom rule

7ef0904

This PR adds an experimental way to register a custom rule for if inductor should partition the graph around an operator. Test Plan: - new test [ghstack-poisoned]

pytorch-bot bot added ciflow/inductor module: inductor labels Sep 19, 2025

zou3519 requested review from BoyuanFeng and eellison September 19, 2025 12:39

ProExpertProg approved these changes Sep 19, 2025

View reviewed changes

BoyuanFeng approved these changes Sep 19, 2025

View reviewed changes

zou3519 added ciflow/trunk Trigger trunk jobs on your pull request topic: not user facing topic category labels Sep 19, 2025

eellison approved these changes Sep 19, 2025

View reviewed changes

zou3519 added this to the 2.9.0 milestone Sep 19, 2025

pytorchmergebot added the merging label Sep 19, 2025

pytorchmergebot closed this in ee7bdd8 Sep 19, 2025

pytorchmergebot added Merged and removed merging labels Sep 19, 2025

This was referenced Sep 20, 2025

[graph partition] Add way to register custom rule (#163310) #163395

Merged

[v.2.9.0] Release Tracker #162497

Closed

ProExpertProg mentioned this pull request Sep 25, 2025

[Feature]: Inductor partitioning should decide what ops to partition on dynamically vllm-project/vllm#25691

Closed

1 task

eellison mentioned this pull request Oct 20, 2025

inductor graph partition x AOTAutograd Cache issue #165471

Open

github-actions bot deleted the gh/zou3519/1201/head branch October 20, 2025 02:17

		_P = ParamSpec("_P")


		_custom_should_partition_fns: weakref.WeakKeyDictionary[

[graph partition] Add way to register custom rule #163310

[graph partition] Add way to register custom rule #163310

Uh oh!

Conversation

zou3519 commented Sep 19, 2025 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Sep 19, 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/163310

⏳ No Failures, 1 Pending

Uh oh!

ProExpertProg left a comment

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!

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!

Choose a reason for hiding this comment

Uh oh!

zou3519 commented Sep 19, 2025

Uh oh!

pytorchmergebot commented Sep 19, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

zou3519 commented Sep 19, 2025 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Sep 19, 2025 •

edited

Loading