[inductor] Multi-output custom op ExternKernelOut lowering by tianrengao · Pull Request #176117 · pytorch/pytorch

tianrengao · 2026-03-02T01:05:56Z

Stack from ghstack (oldest at bottom):

-> [inductor] Multi-output custom op ExternKernelOut lowering #176117
[inductor] Add out-variant discovery + custom op lowering entry point #175116

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

[ghstack-poisoned]

pytorch-bot · 2026-03-02T01:06:00Z

🔗 Helpful Links

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

📄 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 d8993a9 with merge base f0bae19 ():

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

inductor / inductor-cpu-test / test (cpu_inductor_torchbench, 1, 2, linux.2xlarge.amx, unstable) (gh) (#174929)
detectron2_maskrcnn_r_50_fpn

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

pytorch-bot · 2026-03-02T01:06:03Z

This PR needs a `release notes:` label

If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

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

eellison

would you mind combining the above test coverage pr with this pr

Switch the scaled_fp4_quant Python wrapper from calling .out (with pre-allocated buffers) to the functional variant (C++ allocates internally). This enables PyTorch Inductor (2.12+) to automatically discover the .out overload via to_out_variant() and manage output buffers for buffer reuse optimization. Depends on: - pytorch/pytorch#175116 (custom op out-variant lowering hook) - pytorch/pytorch#176117 (multi-output support) Signed-off-by: tianrengao <terrygao87@gmail.com>

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

Extends PR1 single-output lowering to handle multi-output functional custom ops. Adds CustomOpMultiOutputNode (packed node that emits the .out() call) and AllocatingMultiOutput (MultiOutput child with should_allocate=True for buffer reuse). Out arg names from the schema (e.g. out0, out1, output, output_scale) are used in codegen. ghstack-source-id: a7014a2 Pull Request resolved: #176117

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

Extends PR1 single-output lowering to handle multi-output functional custom ops. Adds CustomOpMultiOutputNode (packed node that emits the .out() call) and AllocatingMultiOutput (MultiOutput child with should_allocate=True for buffer reuse). Out arg names from the schema (e.g. out0, out1, output, output_scale) are used in codegen. ghstack-source-id: a8748a4 Pull Request resolved: #176117

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

Extends PR1 single-output lowering to handle multi-output functional custom ops. Adds CustomOpMultiOutputNode (packed node that emits the .out() call) and AllocatingMultiOutput (MultiOutput child with should_allocate=True for buffer reuse). Out arg names from the schema (e.g. out0, out1, output, output_scale) are used in codegen. ghstack-source-id: 169dc39 Pull Request resolved: #176117

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

Extends PR1 single-output lowering to handle multi-output functional custom ops. Adds CustomOpMultiOutputNode (packed node that emits the .out() call) and AllocatingMultiOutput (MultiOutput child with should_allocate=True for buffer reuse). Out arg names from the schema (e.g. out0, out1, output, output_scale) are used in codegen. ghstack-source-id: 0e109f6 Pull Request resolved: #176117

eellison

couple comments.

eellison · 2026-03-04T20:01:50Z

+
+def _codegen_input_args(node: ir.ExternKernel) -> list[str]:
+    """Codegen positional + constant args, excluding out args."""
+    assert ir.is_node_sequence(node.inputs)
+    args = [x.codegen_reference() for x in node.inputs]  # type: ignore[union-attr]
+    for const in node.constant_args:
+        args.append(V.graph.wrapper_code.val_to_arg_str(const))
+    return args
+
+
+def _codegen_kwargs(node: ir.ExternKernel, skip_names: OrderedSet[str]) -> list[str]:
+    """Codegen keyword args, skipping out arg names (appended separately)."""
+    result = []
+    for k, v in node.kwargs.items():
+        if k not in skip_names:
+            result.append(f"{k}={V.graph.wrapper_code.val_to_arg_str(v)}")
+    return result


Can we normalize this, with the rest of the codebase ? see:

pytorch/torch/_inductor/codegen/wrapper.py

Line 584 in c95cc4d

class ExternKernelOutLine(WrapperLine):

this is in a completely different place.

Done. Added ExternKernelMultiOutLine(WrapperLine) right after ExternKernelOutLine in wrapper.py and removed the entire custom_op_out_lowering.py

eellison · 2026-03-04T20:04:04Z

+        return torch.ops.mylib.split_add, torch.ops.mylib.split_add.out
+
+    @parametrize("device", DEVICES)
+    def test_multi_output_lowered_to_out(self, device):


Does this work with cpp wrapper ? can we either add support for this or fallback if its not supported ?

Added cpp_wrapper guard. multioutput .out lowering is skipped under cpp wrapper, falling back to standard FallbackKernel codegen.

cpp wrapper support can be added later when needed.

eellison · 2026-03-04T20:06:20Z

+def try_lower_multi_output_to_out_variant(
+    kernel: OpOverload,
+    example_output: Any,
+    packed: ir.FallbackKernel,
+    *,
+    has_unaligned_input: bool = False,
+) -> Optional[Sequence[AllocatingMultiOutput]]:


if it doesn't make sense to refactor ExternKernelOut to work with the new changes (it may not,)can you add a new class ExternKernelMultiOut or something? and normalize that to where ExternKernel is in the codebase.

Yes, ExternKernelOut can't be extended for multi-output because it's a single-buffer node.

Added ExternKernelMultiOut in ir.py. It inherits FallbackKernel to reuse codegen_args/unflatten_args infrastructure, and overrides codegen() to emit .out() calls. Also deleted custom_op_out_lowering.py

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

Extends PR1 single-output lowering to handle multi-output functional custom ops. Adds CustomOpMultiOutputNode (packed node that emits the .out() call) and AllocatingMultiOutput (MultiOutput child with should_allocate=True for buffer reuse). Out arg names from the schema (e.g. out0, out1, output, output_scale) are used in codegen. ghstack-source-id: 70ba508 Pull Request resolved: #176117

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

Extends PR1 single-output lowering to handle multi-output functional custom ops. Adds CustomOpMultiOutputNode (packed node that emits the .out() call) and AllocatingMultiOutput (MultiOutput child with should_allocate=True for buffer reuse). Out arg names from the schema (e.g. out0, out1, output, output_scale) are used in codegen. ghstack-source-id: 95cbcd6 Pull Request resolved: #176117

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

Extends PR1 single-output lowering to handle multi-output functional custom ops. Adds CustomOpMultiOutputNode (packed node that emits the .out() call) and AllocatingMultiOutput (MultiOutput child with should_allocate=True for buffer reuse). Out arg names from the schema (e.g. out0, out1, output, output_scale) are used in codegen. ghstack-source-id: 61e5d34 Pull Request resolved: #176117

tianrengao · 2026-03-05T17:39:11Z

@pytorchbot merge -i

pytorchmergebot · 2026-03-05T17:49:01Z

Merge started

Your change will be merged while ignoring the following 1 checks: inductor / inductor-cpu-test / test (cpu_inductor_torchbench, 1, 2, linux.2xlarge.amx, unstable)

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

Extends PR1 single-output lowering to handle multi-output functional custom ops. Adds CustomOpMultiOutputNode (packed node that emits the .out() call) and AllocatingMultiOutput (MultiOutput child with should_allocate=True for buffer reuse). Out arg names from the schema (e.g. out0, out1, output, output_scale) are used in codegen. ghstack-source-id: 0e1d843 Pull Request resolved: pytorch/pytorch#176117

The out_variant tag is not yet available in torch==2.10.0. See pytorch/pytorch#176117. Signed-off-by: tianrengao <terrygao87@gmail.com>

…76117) pytorch#175116 Pull Request resolved: pytorch#176117 Approved by: https://github.com/eellison

[inductor] Multi-output custom op ExternKernelOut lowering

1b87153

[ghstack-poisoned]

pytorch-bot Bot added ciflow/inductor module: inductor labels Mar 2, 2026

tianrengao mentioned this pull request Feb 27, 2026

[inductor] Add out-variant discovery + custom op lowering entry point #175116

Closed

tianrengao mentioned this pull request Mar 2, 2026

[inductor] Test coverage for custom op out-variant lowering #176118

Closed

tianrengao added 2 commits March 1, 2026 20:42

Update on "[inductor] Multi-output custom op ExternKernelOut lowering"

df02516

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

Update on "[inductor] Multi-output custom op ExternKernelOut lowering"

742df65

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

eellison reviewed Mar 2, 2026

View reviewed changes

Comment thread torch/_inductor/custom_op_out_lowering.py Outdated

Comment thread torch/_inductor/custom_op_out_lowering.py Outdated

tianrengao mentioned this pull request Mar 3, 2026

[Custom Ops] Add functional + out variant for scaled_fp4_quant vllm-project/vllm#34389

Merged

5 tasks

Update on "[inductor] Multi-output custom op ExternKernelOut lowering"

a68396e

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

Update on "[inductor] Multi-output custom op ExternKernelOut lowering"

7f8dd51

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

Update on "[inductor] Multi-output custom op ExternKernelOut lowering"

25b9b00

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

Update on "[inductor] Multi-output custom op ExternKernelOut lowering"

14c98db

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

tianrengao requested review from drisspg and eellison March 4, 2026 06:59

tianrengao added the release notes: inductor label Mar 4, 2026

eellison reviewed Mar 4, 2026

View reviewed changes

Update on "[inductor] Multi-output custom op ExternKernelOut lowering"

9e8f0d3

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

Update on "[inductor] Multi-output custom op ExternKernelOut lowering"

1a73e13

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

Update on "[inductor] Multi-output custom op ExternKernelOut lowering"

d8993a9

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

eellison approved these changes Mar 5, 2026

View reviewed changes

pytorch-bot Bot added the ciflow/trunk Trigger trunk jobs on your pull request label Mar 5, 2026

pytorchmergebot added the merging label Mar 5, 2026

pytorchmergebot added the Merged label Mar 5, 2026

pytorchmergebot closed this in df4c495 Mar 5, 2026

pytorchmergebot removed the merging label Mar 5, 2026

tianrengao added a commit to tianrengao/vllm that referenced this pull request Mar 13, 2026

Remove at::Tag::out_variant pending PyTorch upgrade

01736ec

The out_variant tag is not yet available in torch==2.10.0. See pytorch/pytorch#176117. Signed-off-by: tianrengao <terrygao87@gmail.com>

EmanueleCoradin pushed a commit to EmanueleCoradin/pytorch that referenced this pull request Mar 30, 2026

[inductor] Multi-output custom op ExternKernelOut lowering (pytorch#1…

eb0d964

…76117) pytorch#175116 Pull Request resolved: pytorch#176117 Approved by: https://github.com/eellison

github-actions Bot deleted the gh/tianrengao/29/head branch April 5, 2026 02:26

Conversation

tianrengao commented Mar 2, 2026 • 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 Mar 2, 2026 • 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/176117

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

pytorch-bot Bot commented Mar 2, 2026

This PR needs a release notes: label

Uh oh!

eellison left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

eellison left a comment

Choose a reason for hiding this comment

Uh oh!

eellison Mar 4, 2026

Choose a reason for hiding this comment

Uh oh!

tianrengao Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

eellison Mar 4, 2026

Choose a reason for hiding this comment

Uh oh!

tianrengao Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

eellison Mar 4, 2026

Choose a reason for hiding this comment

Uh oh!

tianrengao Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

tianrengao commented Mar 5, 2026

Uh oh!

pytorchmergebot commented Mar 5, 2026

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

tianrengao commented Mar 2, 2026 •

edited by pytorch-bot Bot

Loading

pytorch-bot Bot commented Mar 2, 2026 •

edited

Loading

This PR needs a `release notes:` label