Fix no source name in backward kernel names; Add flex_attention HOP to "original_aten" node meta by yushangdi · Pull Request #167749 · pytorch/pytorch

yushangdi · 2025-11-13T18:48:05Z

Add torch.fx.experimental.proxy_tensor.set_original_aten_op() around flex_atention HOP dispatch so we have original_aten populated for flex_attention
Update the usages of original_aten to also expect HOP in addition to OpOverload

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

pytorch-bot · 2025-11-13T18:48:12Z

🔗 Helpful Links

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

📄 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! (4 Unrelated Failures)

As of commit 4e35032 with merge base a76dd6b ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

inductor / inductor-cpu-test / test (cpu_inductor_torchbench, 2, 2, linux.2xlarge.amx) (gh) (similar failure)
stable_diffusion_unet
inductor / inductor-cpu-test / test (dynamic_cpu_inductor_torchbench, 2, 2, linux.2xlarge.amx) (gh) (similar failure)
stable_diffusion_unet
inductor / inductor-test / test (inductor_torchbench, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
stable_diffusion_unet

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)
backends/xnnpack/test/recipes/test_xnnpack_recipes.py::TestXnnpackRecipes::test_int8_static_quant_recipe

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

pytorch-bot · 2025-11-13T18:48:13Z

The label module: inductor is only applicable to issues and has been removed. Please only use this label on issues.

pytorch-bot · 2025-11-13T18:48:25Z

The label module: inductor is only applicable to issues and has been removed. Please only use this label on issues.

pytorch-bot · 2025-11-13T19:33:12Z

The label module: inductor is only applicable to issues and has been removed. Please only use this label on issues.

pytorch-bot · 2025-11-13T19:33:13Z

The label module: dynamo is only applicable to issues and has been removed. Please only use this label on issues.

pytorch-bot · 2025-11-13T19:33:20Z

The label module: dynamo is only applicable to issues and has been removed. Please only use this label on issues.

pytorch-bot · 2025-11-13T19:33:20Z

The label module: inductor is only applicable to issues and has been removed. Please only use this label on issues.

pytorch-bot · 2025-11-13T19:37:42Z

The label module: inductor is only applicable to issues and has been removed. Please only use this label on issues.

pytorch-bot · 2025-11-13T19:37:52Z

The label module: dynamo is only applicable to issues and has been removed. Please only use this label on issues.

pytorch-bot · 2025-11-13T19:37:52Z

The label module: inductor is only applicable to issues and has been removed. Please only use this label on issues.

drisspg · 2025-11-13T19:40:48Z

what do we see now? also does this apply to the generated kernels?

yushangdi · 2025-11-13T20:01:32Z

what do we see now? also does this apply to the generated kernels?

example output: triton_tem_fused_flex_attention_backward_zeros_1

V1113 11:32:05.735000 42282 /data/users/shangdiy/pytorch/torch/_inductor/graph.py:2431] [1/0] [__output_code]             # Topologically Sorted Source Nodes: [full_default_4, flex_attention_backward], Original ATen: [aten.zeros, flex_attention_backward]
V1113 11:32:05.735000 42282 /data/users/shangdiy/pytorch/torch/_inductor/graph.py:2431] [1/0] [__output_code]             stream0 = get_raw_stream(0)
V1113 11:32:05.735000 42282 /data/users/shangdiy/pytorch/torch/_inductor/graph.py:2431] [1/0] [__output_code]             triton_tem_fused_flex_attention_backward_zeros_1.run(primals_2, primals_5, primals_6, getitem_3, buf1, tangents_1, buf3, buf4, full, full_default, convert_element_type, convert_element_type_1, buf5, buf6, buf7, buf8, buf9, s50, s41, s0, ((15 + s0) // 16) + ((15 + s50) // 16), 1, 2, stream=stream0)

pytorch-bot · 2025-11-13T22:32:24Z

The label module: dynamo is only applicable to issues and has been removed. Please only use this label on issues.

pytorch-bot · 2025-11-13T22:32:24Z

The label module: inductor is only applicable to issues and has been removed. Please only use this label on issues.

yushangdi · 2025-11-14T02:15:41Z

@pytorchbot merge

pytorchmergebot · 2025-11-14T02:18:10Z

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

…o "original_aten" node meta (pytorch#167749) Fixes pytorch#167706 - Add `torch.fx.experimental.proxy_tensor.set_original_aten_op()` around flex_atention HOP dispatch so we have `original_aten` populated for flex_attention - Update the usages of `original_aten` to also expect HOP in addition to OpOverload Pull Request resolved: pytorch#167749 Approved by: https://github.com/drisspg

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

pytorch-bot bot removed the module: inductor label Nov 13, 2025

pytorch-bot bot added module: inductor and removed module: inductor labels Nov 13, 2025

yushangdi mentioned this pull request Nov 13, 2025

Improve Inductor generated kernel names: important pieces like flex attention backward not included #167706

Closed

yushangdi force-pushed the sy_kernel_bw_name branch from 45b1d8f to 2fe38de Compare November 13, 2025 19:33

yushangdi requested review from drisspg and zou3519 as code owners November 13, 2025 19:33

pytorch-bot bot added module: dynamo module: inductor release notes: fx release notes category and removed module: inductor module: dynamo labels Nov 13, 2025

pytorch-bot bot added module: dynamo module: inductor and removed module: dynamo module: inductor labels Nov 13, 2025

yushangdi changed the title ~~Fix no source name in backward kernel names~~ Fix no source name in backward kernel names; Add flex_attention HOP to "original_aten" Nov 13, 2025

pytorch-bot bot added module: dynamo module: inductor labels Nov 13, 2025

yushangdi changed the title ~~Fix no source name in backward kernel names; Add flex_attention HOP to "original_aten"~~ Fix no source name in backward kernel names; Add flex_attention HOP to "original_aten" node meta Nov 13, 2025

pytorch-bot bot removed the module: inductor label Nov 13, 2025

pytorch-bot bot added module: dynamo module: inductor and removed module: dynamo module: inductor labels Nov 13, 2025

yushangdi requested a review from ezyang November 13, 2025 20:29

Fix no source name in backward kernel names

4e35032

yushangdi requested a review from angelayi November 13, 2025 22:32

yushangdi added the ciflow/trunk Trigger trunk jobs on your pull request label Nov 13, 2025

yushangdi force-pushed the sy_kernel_bw_name branch from 2fe38de to 4e35032 Compare November 13, 2025 22:32

pytorch-bot bot added module: dynamo module: inductor and removed module: dynamo module: inductor labels Nov 13, 2025

drisspg approved these changes Nov 13, 2025

View reviewed changes

pytorchmergebot added the merging label Nov 14, 2025

pytorchmergebot added the Merged label Nov 14, 2025

pytorchmergebot closed this in 79317dc Nov 14, 2025

pytorchmergebot removed the merging label Nov 14, 2025

github-actions bot deleted the sy_kernel_bw_name branch December 15, 2025 02:21

Conversation

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

Uh oh!

Uh oh!

pytorch-bot bot commented Nov 13, 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/167749

✅ You can merge normally! (4 Unrelated Failures)

Uh oh!

pytorch-bot bot commented Nov 13, 2025

Uh oh!

pytorch-bot bot commented Nov 13, 2025

Uh oh!

pytorch-bot bot commented Nov 13, 2025

Uh oh!

pytorch-bot bot commented Nov 13, 2025

Uh oh!

pytorch-bot bot commented Nov 13, 2025

Uh oh!

pytorch-bot bot commented Nov 13, 2025

Uh oh!

pytorch-bot bot commented Nov 13, 2025

Uh oh!

pytorch-bot bot commented Nov 13, 2025

Uh oh!

pytorch-bot bot commented Nov 13, 2025

Uh oh!

drisspg commented Nov 13, 2025

Uh oh!

yushangdi commented Nov 13, 2025

Uh oh!

pytorch-bot bot commented Nov 13, 2025

Uh oh!

pytorch-bot bot commented Nov 13, 2025

Uh oh!

yushangdi commented Nov 14, 2025

Uh oh!

pytorchmergebot commented Nov 14, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yushangdi commented Nov 13, 2025 •

edited

Loading

pytorch-bot bot commented Nov 13, 2025 •

edited

Loading