[AOTAutograd] Fix static_input_indices not offset when effect tokens are prepended by wmhst7 · Pull Request #175904 · pytorch/pytorch

wmhst7 · 2026-02-26T22:47:40Z

Summary

When effectful ops (e.g., with_effects) are present, handle_effect_tokens_fn() prepends effect token placeholders to the input args. However, static_input_indices in ViewAndMutationMeta is computed before this prepending and is not adjusted afterwards. This causes indices to point to wrong inputs, leading to issues like unnecessary CUDA graph re-recording.

Problem

In handle_effect_tokens_fn(), effect tokens are prepended to args:

additional_fwd_token_inputs = [torch.tensor([])] * num_tokens
args = [*additional_fwd_token_inputs, *args]  # tokens prepended at index 0

But meta.static_input_indices is not offset by num_tokens. When these indices are later used (e.g., by CUDAGraph's check_invariants), they point to the wrong inputs:

Before tokens: args=[activation, weight], static_input_indices=[1] → weight ✓
After tokens: args=[token, activation, weight], static_input_indices=[1] → activation ✗
Expected: static_input_indices=[2] (offset by num_tokens=1) → weight ✓

Impact

Activations get incorrectly marked as static inputs
CUDAGraph's check_invariants sees data pointer changes for "static" inputs
This triggers unnecessary re-recording, causing performance degradation

Fix

Offset static_input_indices by num_tokens after prepending effect tokens in the forward-only (trace_joint=False) path:

  if num_tokens > 0:
      meta.static_input_indices = [
          idx + num_tokens for idx in meta.static_input_indices
      ]

Unit Test

Added test_static_input_indices_with_effect_tokens in test/functorch/test_aotdispatch.py which:

Registers a custom effectful op via _register_effectful_op
Compiles a function with torch.compile using a metadata-capturing backend
Verifies that all static_input_indices are >= num_tokens after effect tokens are prepended (i.e., no index
incorrectly points to a token input)

cc @yanboliang

pytorch-bot · 2026-02-26T22:47:45Z

🔗 Helpful Links

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

📄 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 36ac958 with merge base ca7ffb7 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

trunk / linux-jammy-cuda13.0-py3.10-gcc11 / test (default, 4, 5, lf.linux.g6.4xlarge.experimental.nvidia.gpu) (gh) (similar failure)
test/dynamo/test_misc.py::MiscTests::test_assume_32_bit_indexing

BROKEN TRUNK - The following jobs failed but was present on the merge base:

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

trunk / linux-jammy-cuda13.0-py3.10-gcc11 / test (default, 2, 5, lf.linux.g6.4xlarge.experimental.nvidia.gpu) (gh) (trunk failure)
test/dynamo/test_dynamic_shapes.py::DynamicShapesMiscTests::test_assume_32_bit_indexing_dynamic_shapes
trunk / linux-jammy-rocm-py3.10 / test (default, 2, 6, linux.rocm.gpu.gfx950.1) (gh) (trunk failure)
test/dynamo/test_misc.py::MiscTests::test_assume_32_bit_indexing
trunk / linux-jammy-rocm-py3.10 / test (default, 3, 6, linux.rocm.gpu.gfx950.1) (gh) (trunk failure)
test/dynamo/test_dynamic_shapes.py::DynamicShapesMiscTests::test_assume_32_bit_indexing_dynamic_shapes

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

pytorch-bot · 2026-02-26T22:47:48Z

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.

linux-foundation-easycla · 2026-03-03T05:22:46Z

The committers listed above are authorized under a signed CLA.

✅ login: wmhst7 / name: Mingheng Wu (7dc3a35, 1a49e8e, 36ac958)

yanboliang · 2026-03-03T22:20:02Z

@zou3519 @angelayi

angelayi

thanks!

Fix line-too-long lint error in graph_capture_wrappers.py and add a test verifying that static_input_indices are correctly offset when effect tokens are prepended to inputs.

wmhst7 · 2026-03-05T04:54:13Z

@pytorchbot merge

pytorchmergebot · 2026-03-05T04:56:20Z

Merge failed

Reason: 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.

Details for Dev Infra team

Raised by workflow job

pytorch-bot · 2026-03-05T04:56:25Z

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.

wmhst7 · 2026-03-05T04:57:24Z

@pytorchbot label "topic: not user facing"

wmhst7 · 2026-03-05T04:58:08Z

@pytorchbot merge

pytorchmergebot · 2026-03-05T05:00:16Z

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

Update graph_capture_wrappers.py

7dc3a35

wmhst7 requested review from aorenste and bdhirsh as code owners February 26, 2026 22:47

pytorchbot added the open source label Feb 26, 2026

soulitzer added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Feb 27, 2026

wmhst7 requested review from Chillee and ezyang as code owners March 3, 2026 05:22

wmhst7 force-pushed the wmhst7-patch-1 branch from 1943988 to ebc7be2 Compare March 3, 2026 05:37

angelayi approved these changes Mar 3, 2026

View reviewed changes

Fix lint and add test for static_input_indices effect token offset

1a49e8e

Fix line-too-long lint error in graph_capture_wrappers.py and add a test verifying that static_input_indices are correctly offset when effect tokens are prepended to inputs.

wmhst7 force-pushed the wmhst7-patch-1 branch from ebc7be2 to 1a49e8e Compare March 4, 2026 00:36

Fix lint: collapse custom_op decorator to single line

36ac958

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 removed the merging label Mar 5, 2026

pytorch-bot bot added the topic: not user facing topic category label Mar 5, 2026

pytorchmergebot added the merging label Mar 5, 2026

pytorchmergebot added the Merged label Mar 5, 2026

pytorchmergebot closed this in 36ed9aa Mar 5, 2026

pytorchmergebot removed the merging label Mar 5, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[AOTAutograd] Fix static_input_indices not offset when effect tokens are prepended#175904

[AOTAutograd] Fix static_input_indices not offset when effect tokens are prepended#175904
wmhst7 wants to merge 3 commits intopytorch:mainfrom
wmhst7:wmhst7-patch-1

wmhst7 commented Feb 26, 2026 •

edited

Loading

Uh oh!

pytorch-bot bot commented Feb 26, 2026 •

edited

Loading

Uh oh!

pytorch-bot bot commented Feb 26, 2026

Uh oh!

linux-foundation-easycla bot commented Mar 3, 2026 •

edited

Loading

Uh oh!

yanboliang commented Mar 3, 2026

Uh oh!

angelayi left a comment

Uh oh!

wmhst7 commented Mar 5, 2026

Uh oh!

pytorchmergebot commented Mar 5, 2026

Uh oh!

pytorch-bot bot commented Mar 5, 2026

Uh oh!

wmhst7 commented Mar 5, 2026

Uh oh!

wmhst7 commented Mar 5, 2026

Uh oh!

pytorchmergebot commented Mar 5, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

wmhst7 commented Feb 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Problem

Impact

Fix

Unit Test

Uh oh!

pytorch-bot bot commented Feb 26, 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/175904

✅ You can merge normally! (4 Unrelated Failures)

Uh oh!

pytorch-bot bot commented Feb 26, 2026

This PR needs a release notes: label

Uh oh!

linux-foundation-easycla bot commented Mar 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yanboliang commented Mar 3, 2026

Uh oh!

angelayi left a comment

Choose a reason for hiding this comment

Uh oh!

wmhst7 commented Mar 5, 2026

Uh oh!

pytorchmergebot commented Mar 5, 2026

Merge failed

Uh oh!

pytorch-bot bot commented Mar 5, 2026

This PR needs a release notes: label

Uh oh!

wmhst7 commented Mar 5, 2026

Uh oh!

wmhst7 commented Mar 5, 2026

Uh oh!

pytorchmergebot commented Mar 5, 2026

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

wmhst7 commented Feb 26, 2026 •

edited

Loading

pytorch-bot bot commented Feb 26, 2026 •

edited

Loading

This PR needs a `release notes:` label

linux-foundation-easycla bot commented Mar 3, 2026 •

edited

Loading

This PR needs a `release notes:` label