[inductor] remove redundant view/permute earlier by shunting314 · Pull Request #170649 · pytorch/pytorch

shunting314 · 2025-12-17T07:51:35Z

Stack from ghstack (oldest at bottom):

2 benefits

remove these redundancy earlier so that AutoChunker does not need to handle redundant view pairs
some attention patterns add redudant views in the pattern on purpose for matching against the compiled graph. By removing these redundancy in the compiled graph earlier, we can simplify those patterns a bit by not injecting redundant views.

test by the existing tests in test_fused_attention.py

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

[ghstack-poisoned]

pytorch-bot · 2025-12-17T07:51:38Z

🔗 Helpful Links

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

📄 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 b7a971d with merge base 2708634 ():

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

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

inductor / unit-test / inductor-pallas-cpu-test / test (inductor-pallas-cpu, 1, 1, linux.12xlarge) (gh) (trunk failure)
Process completed with exit code 1.

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

eellison

Should we apply these canonicalization patterns to the patterns generated from register_replacement, as well ?

eellison · 2025-12-17T18:20:54Z

-        score.size(0), score.size(1), score.size(2), score.size(3)
-    )
-    return viewd_score2.float().softmax(dim=-1).type_as(query).matmul(value), key, value
+    return score.float().softmax(dim=-1).type_as(query).matmul(value), key, value


ideally we wouldn't need to change these

shunting314 · 2025-12-17T19:17:19Z

Should we apply these canonicalization patterns to the patterns generated from register_replacement, as well ?

I think both should work. But the change to the attn pattern makes them simpler. So I would still prefer changing those patterns.

On the other hand, we can apply those canonicalization patterns in the trace function generating patterns as the same time. But I have 2 concerns

applying other patterns when generating a more complex pattern seems overall tricky. Maybe it's fine if the applied patterns are simple enough
I think it's more helpful if the canonicalization patterns removes redundancies generated due to tracing (rather than due to how the code being traced is written). If we can improve the code being traced to remove those redundancy, I like better to just change the original code.

I'm open to more discussions.

2 benefits 1. remove these redundancy earlier so that AutoChunker does not need to handle redundant view pairs 2. some attention patterns add redudant views in the pattern on purpose for matching against the compiled graph. By removing these redundancy in the compiled graph earlier, we can simplify those patterns a bit by not injecting redundant views. test by the existing tests in test_fused_attention.py cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo [ghstack-poisoned]

eellison

Since this runs prior to all post grad, and joint graph passes, can we apply this to user patterns as well ? this would just be added as a pass to register_replacement.

2 benefits 1. remove these redundancy earlier so that AutoChunker does not need to handle redundant view pairs 2. some attention patterns add redudant views in the pattern on purpose for matching against the compiled graph. By removing these redundancy in the compiled graph earlier, we can simplify those patterns a bit by not injecting redundant views. test by the existing tests in test_fused_attention.py cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo [ghstack-poisoned]

pytorchmergebot · 2025-12-23T19:29:11Z

Starting merge as part of PR stack under #136702

2 benefits 1. remove these redundancy earlier so that AutoChunker does not need to handle redundant view pairs 2. some attention patterns add redudant views in the pattern on purpose for matching against the compiled graph. By removing these redundancy in the compiled graph earlier, we can simplify those patterns a bit by not injecting redundant views. test by the existing tests in test_fused_attention.py cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo [ghstack-poisoned]

pytorchmergebot · 2025-12-24T05:23:51Z

Starting merge as part of PR stack under #136702

The AutoChunker defines the following chunking metadata and propagate them thru the subgraphs that we chunk 1. scale_by: The AutoChunker is only enabled if there is a single scalar tangent. To decouple the dependency on tangent for the bwd subgraph so we can compute them in fwd, we pretend the tangent is 1 first and record it in 'scale_by' . This metadata get propagated and when we cancel the chunking effect in the end of bwd subgraph, we apply the scaling. 2. chunk_dim: record which dimension of the tensor get chunked 3. need_sum: if true, the original Tensor is the sum (rather than concat) of each chunked tensors. One important implementation detail is, we need put chunked subgraph in a HOP (use invoke_subgraph here). Otherwise Inductor fuse across these subgraphs and results in no peak memory saving. Here are some early benchmarking result on GPT2. - 64 chunks: - final 19 iters avg: 242.550ms - peak memory consumption: 12603 MiB - 32 chunks: - final 19 iters avg: 206.180ms - peak memory consumption: 12880 MiB - 16 chunks - final 19 iters avg: 196.997ms - peak memory consumption: 13267 MiB - 8 chunks - final 19 iters avg: 194.924ms - peak memory consumption: 14049 MiB With 64 chunks, our peak memory is smaller than llm.c's 13.4GB. I also tried the AutoChunker on PT2 OSS benchmarks to verify the numerical. By default our accuracy test picks a very small batch size. This makes AutoChunker get skipped. I force batch_size to be 16 for BertForMaskedLM to trigger the AutoChunker and verified the numerical correctness. Pull Request resolved: #136702 Approved by: https://github.com/jansel ghstack dependencies: #170649

huydhn · 2025-12-24T10:07:38Z

@pytorchbot revert -m 'Sorry for reverting your change but this seems to break vLLM CI' -c nosignal

GH job link HUD commit link

pytorchmergebot · 2025-12-24T10:09:22Z

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

This reverts commit 8e09f22. Reverted #170649 on behalf of https://github.com/huydhn due to Sorry for reverting your change but this seems to break vLLM CI ([comment](#170649 (comment)))

pytorchmergebot · 2025-12-24T10:09:32Z

@shunting314 your PR has been successfully reverted.

shunting314 · 2025-12-26T22:23:58Z

Since this runs prior to all post grad, and joint graph passes, can we apply this to user patterns as well ? this would just be added as a pass to register_replacement.

Applying these to user patterns (in fwd_only) cause duplicate patterns registered from vllm. If we really want to do this, we'd need to change vllm first. Repro for reference:

time pytest -v -s tests/compile/test_pass_manager.py

vllm pinned commit: 7c73ceb5812ace65e0d1b6ada3622b8b9f0400c0

2 benefits 1. remove these redundancy earlier so that AutoChunker does not need to handle redundant view pairs 2. some attention patterns add redudant views in the pattern on purpose for matching against the compiled graph. By removing these redundancy in the compiled graph earlier, we can simplify those patterns a bit by not injecting redundant views. test by the existing tests in test_fused_attention.py cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo [ghstack-poisoned]

pytorchmergebot · 2025-12-27T01:56:39Z

Starting merge as part of PR stack under #136702

The AutoChunker defines the following chunking metadata and propagate them thru the subgraphs that we chunk 1. scale_by: The AutoChunker is only enabled if there is a single scalar tangent. To decouple the dependency on tangent for the bwd subgraph so we can compute them in fwd, we pretend the tangent is 1 first and record it in 'scale_by' . This metadata get propagated and when we cancel the chunking effect in the end of bwd subgraph, we apply the scaling. 2. chunk_dim: record which dimension of the tensor get chunked 3. need_sum: if true, the original Tensor is the sum (rather than concat) of each chunked tensors. One important implementation detail is, we need put chunked subgraph in a HOP (use invoke_subgraph here). Otherwise Inductor fuse across these subgraphs and results in no peak memory saving. Here are some early benchmarking result on GPT2. - 64 chunks: - final 19 iters avg: 242.550ms - peak memory consumption: 12603 MiB - 32 chunks: - final 19 iters avg: 206.180ms - peak memory consumption: 12880 MiB - 16 chunks - final 19 iters avg: 196.997ms - peak memory consumption: 13267 MiB - 8 chunks - final 19 iters avg: 194.924ms - peak memory consumption: 14049 MiB With 64 chunks, our peak memory is smaller than llm.c's 13.4GB. I also tried the AutoChunker on PT2 OSS benchmarks to verify the numerical. By default our accuracy test picks a very small batch size. This makes AutoChunker get skipped. I force batch_size to be 16 for BertForMaskedLM to trigger the AutoChunker and verified the numerical correctness. Pull Request resolved: #136702 Approved by: https://github.com/jansel ghstack dependencies: #170649

2 benefits 1. remove these redundancy earlier so that AutoChunker does not need to handle redundant view pairs 2. some attention patterns add redudant views in the pattern on purpose for matching against the compiled graph. By removing these redundancy in the compiled graph earlier, we can simplify those patterns a bit by not injecting redundant views. test by the existing tests in test_fused_attention.py Pull Request resolved: pytorch#170649 Approved by: https://github.com/eellison, https://github.com/jansel

The AutoChunker defines the following chunking metadata and propagate them thru the subgraphs that we chunk 1. scale_by: The AutoChunker is only enabled if there is a single scalar tangent. To decouple the dependency on tangent for the bwd subgraph so we can compute them in fwd, we pretend the tangent is 1 first and record it in 'scale_by' . This metadata get propagated and when we cancel the chunking effect in the end of bwd subgraph, we apply the scaling. 2. chunk_dim: record which dimension of the tensor get chunked 3. need_sum: if true, the original Tensor is the sum (rather than concat) of each chunked tensors. One important implementation detail is, we need put chunked subgraph in a HOP (use invoke_subgraph here). Otherwise Inductor fuse across these subgraphs and results in no peak memory saving. Here are some early benchmarking result on GPT2. - 64 chunks: - final 19 iters avg: 242.550ms - peak memory consumption: 12603 MiB - 32 chunks: - final 19 iters avg: 206.180ms - peak memory consumption: 12880 MiB - 16 chunks - final 19 iters avg: 196.997ms - peak memory consumption: 13267 MiB - 8 chunks - final 19 iters avg: 194.924ms - peak memory consumption: 14049 MiB With 64 chunks, our peak memory is smaller than llm.c's 13.4GB. I also tried the AutoChunker on PT2 OSS benchmarks to verify the numerical. By default our accuracy test picks a very small batch size. This makes AutoChunker get skipped. I force batch_size to be 16 for BertForMaskedLM to trigger the AutoChunker and verified the numerical correctness. Pull Request resolved: pytorch#136702 Approved by: https://github.com/jansel ghstack dependencies: pytorch#170649

)" This reverts commit 8e09f22. Reverted pytorch#170649 on behalf of https://github.com/huydhn due to Sorry for reverting your change but this seems to break vLLM CI ([comment](pytorch#170649 (comment)))

2 benefits 1. remove these redundancy earlier so that AutoChunker does not need to handle redundant view pairs 2. some attention patterns add redudant views in the pattern on purpose for matching against the compiled graph. By removing these redundancy in the compiled graph earlier, we can simplify those patterns a bit by not injecting redundant views. test by the existing tests in test_fused_attention.py Pull Request resolved: pytorch#170649 Approved by: https://github.com/eellison, https://github.com/jansel

The AutoChunker defines the following chunking metadata and propagate them thru the subgraphs that we chunk 1. scale_by: The AutoChunker is only enabled if there is a single scalar tangent. To decouple the dependency on tangent for the bwd subgraph so we can compute them in fwd, we pretend the tangent is 1 first and record it in 'scale_by' . This metadata get propagated and when we cancel the chunking effect in the end of bwd subgraph, we apply the scaling. 2. chunk_dim: record which dimension of the tensor get chunked 3. need_sum: if true, the original Tensor is the sum (rather than concat) of each chunked tensors. One important implementation detail is, we need put chunked subgraph in a HOP (use invoke_subgraph here). Otherwise Inductor fuse across these subgraphs and results in no peak memory saving. Here are some early benchmarking result on GPT2. - 64 chunks: - final 19 iters avg: 242.550ms - peak memory consumption: 12603 MiB - 32 chunks: - final 19 iters avg: 206.180ms - peak memory consumption: 12880 MiB - 16 chunks - final 19 iters avg: 196.997ms - peak memory consumption: 13267 MiB - 8 chunks - final 19 iters avg: 194.924ms - peak memory consumption: 14049 MiB With 64 chunks, our peak memory is smaller than llm.c's 13.4GB. I also tried the AutoChunker on PT2 OSS benchmarks to verify the numerical. By default our accuracy test picks a very small batch size. This makes AutoChunker get skipped. I force batch_size to be 16 for BertForMaskedLM to trigger the AutoChunker and verified the numerical correctness. Pull Request resolved: pytorch#136702 Approved by: https://github.com/jansel ghstack dependencies: pytorch#170649

ghstack-source-id: 6c50026 Pull Request resolved: pytorch/pytorch#170649

[inductor] remove redundant view/permute earlier

3b8d307

[ghstack-poisoned]

pytorch-bot Bot added ciflow/inductor module: inductor labels Dec 17, 2025

shunting314 mentioned this pull request Dec 17, 2025

[Inductor] auto-chunker #136702

Closed

shunting314 requested review from eellison, jansel and v0i0 December 17, 2025 18:17

eellison reviewed Dec 17, 2025

View reviewed changes

v0i0 reviewed Dec 17, 2025

View reviewed changes

Comment thread torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_21.py

This was referenced Dec 18, 2025

NOT-FOR-LAND: enable auto-chunker by default #170730

Closed

NOT-FOR-LAND: temporary change for torchtitan integration #170830

Closed

eellison reviewed Dec 22, 2025

View reviewed changes

eellison approved these changes Dec 22, 2025

View reviewed changes

jansel approved these changes Dec 23, 2025

View reviewed changes

shunting314 added the topic: not user facing topic category label Dec 23, 2025

pytorchmergebot closed this in 8e09f22 Dec 24, 2025

pytorchmergebot added the Merged label Dec 24, 2025

pytorchmergebot added Reverted ci-no-td Do not run TD on this PR labels Dec 24, 2025

pytorchmergebot reopened this Dec 24, 2025

shunting314 mentioned this pull request Dec 27, 2025

[autochunker] support gradient accumulation #171359

Closed

pytorchmergebot closed this in d50721a Dec 27, 2025

shunting314 mentioned this pull request Dec 29, 2025

[autochunker] override num-chunks #171477

Closed

SergeyTyshkevich pushed a commit to SergeyTyshkevich/chart2 that referenced this pull request Jan 19, 2026

[inductor] remove redundant view/permute earlier

b6c1708

ghstack-source-id: 6c50026 Pull Request resolved: pytorch/pytorch#170649

github-actions Bot deleted the gh/shunting314/274/head branch January 27, 2026 02:19

Navigation Menu

Uh oh!

[inductor] remove redundant view/permute earlier - #170649

[inductor] remove redundant view/permute earlier#170649
shunting314 wants to merge 5 commits into
gh/shunting314/274/basefrom
gh/shunting314/274/head

shunting314 commented Dec 17, 2025 •

edited

Loading

Uh oh!

pytorch-bot Bot commented Dec 17, 2025 •

edited

Loading

Uh oh!

eellison left a comment

Uh oh!

eellison Dec 17, 2025

Uh oh!

shunting314 commented Dec 17, 2025 •

edited

Loading

Uh oh!

Uh oh!

eellison left a comment

Uh oh!

pytorchmergebot commented Dec 23, 2025

Uh oh!

pytorchmergebot commented Dec 24, 2025

Uh oh!

huydhn commented Dec 24, 2025

Uh oh!

pytorchmergebot commented Dec 24, 2025

Uh oh!

pytorchmergebot commented Dec 24, 2025

Uh oh!

shunting314 commented Dec 26, 2025

Uh oh!

pytorchmergebot commented Dec 27, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

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

Uh oh!

Uh oh!

pytorch-bot Bot commented Dec 17, 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/170649

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

eellison left a comment

Choose a reason for hiding this comment

Uh oh!

eellison Dec 17, 2025

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Uh oh!

Uh oh!

eellison left a comment

Choose a reason for hiding this comment

Uh oh!

pytorchmergebot commented Dec 23, 2025

Uh oh!

pytorchmergebot commented Dec 24, 2025

Uh oh!

huydhn commented Dec 24, 2025

Uh oh!

pytorchmergebot commented Dec 24, 2025

Uh oh!

pytorchmergebot commented Dec 24, 2025

Uh oh!

shunting314 commented Dec 26, 2025

Uh oh!

pytorchmergebot commented Dec 27, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

shunting314 commented Dec 17, 2025 •

edited

Loading

pytorch-bot Bot commented Dec 17, 2025 •

edited

Loading

shunting314 commented Dec 17, 2025 •

edited

Loading