AOTInductor Input Mutation

   blog 26 minutes read (About 3955 words)  visits

Introduction

AOTInductor is a PyTorch compiler backend that compiles PyTorch models into optimized inference engines. In my previous impression, I thought that AOTInductor relies on functionalization and does not support memory mutations. Any memory mutations will be a consequence of AOTInductor’s internal optimizations, and the user cannot completely control them. When it comes to inputs, I could not imagine that AOTInductor would allow users to mutate inputs in-place, because it would break the functionalization assumption. However, it turns out that I was wrong. Actually, AOTInductor strictly relies on functionalization, meaning it expects a clean mathematical graph without internal memory mutations or global side effects before executing its code-generation phase. It does not mean that optimized engine produced by AOTInductor code-generation cannot have side effects, such as in-place input mutations.

The motivation of input mutation is the scenario that sometimes we would just like to get an output tensor that only changes a very small fraction of a large input tensor. Out-of-place operations would require allocating a new tensor and copying the unchanged data from the input tensor to the output tensor, which is inefficient. In-place operations can avoid this overhead by directly modifying the input tensor. In this blog post, I will demonstrate how to enable in-place input mutation optimizations in AOTInductor.

AOTInductor Input Mutation

The key of using in-place input mutation operations is to explicitly use in-place mutations, such as x.mul_(2) instead of x.mul(2), for the input tensors in the PyTorch model. For custom Triton kernels, in addition to an implementation of Triton kernel that mutates the input tensor in-place, wrapping it with triton_op and wrap_triton, the mutates_args argument should be used to explicitly indicate which input arguments are mutated.

After exporting the model with torch.export.export(..., strict=True), the ExportedProgram will still have the in-place operation torch.ops.aten.mul_ in the graph, but the top-level graph signature will not have any user_inputs_to_mutate. After decomposing the ExportedProgram with run_decompositions, the decomposed ExportedProgram will become functionalized, and the in-place operation will be replaced with an out-of-place operation torch.ops.aten.mul. The input mutation side effect will then be tracked in the decomposed graph signature, and the input tensor will be listed in user_inputs_to_mutate. The decomposed ExportedProgram, actually as well as the original ExportedProgram, can then be compiled with AOTInductor, and AOTInductor will respect the input mutation side effect and generate an optimized engine that performs in-place input mutation.

aoti_input_mutation_example.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
import os
import tempfile
import torch
from torch.library import triton_op, wrap_triton
import triton
import triton.language as tl


@triton.jit
def _triton_mul2_kernel(
    x_ptr,
    n_elements,
    BLOCK_SIZE: "tl.constexpr",
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    x = x * 2.0
    tl.store(x_ptr + offsets, x, mask=mask)


@triton_op("mylib::triton_mul2_mutation", mutates_args=("x", ))
def triton_mul2_mutation(x: torch.Tensor) -> None:
    """Custom operator wrapping a Triton kernel that mutates x in-place using triton_op and wrap_triton."""
    if x.is_cuda:
        n_elements = x.numel()

        def grid(meta):
            return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), )

        wrap_triton(_triton_mul2_kernel)[grid](x, n_elements, BLOCK_SIZE=1024)
    else:
        x.mul_(2)


class UserInputMutationTriton(torch.nn.Module):
    """Mutates input x in-place using a custom Triton kernel wrapped with triton_op and wrap_triton."""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        triton_mul2_mutation(x)
        return x.cos()


class UserInputMutationNoReturn(torch.nn.Module):
    """Mutates input x in-place without returning x."""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x.mul_(2)
        return x.cos()


class UserInputMutationWithReturn(torch.nn.Module):
    """Mutates input x in-place AND explicitly returns x alongside the result."""

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        x.mul_(2)
        return x, x.cos()


def test_mutation_pipeline(model_class: type[torch.nn.Module],
                           name: str) -> None:
    print(f"\n{'='*75}")
    print(f" TESTING: {name}")
    print(f"{'='*75}")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model_class().to(device).eval()
    example_args = (torch.randn(3, 2, device=device), )

    # --------------------------------------------------------------------------
    # 1. Export with strict=True & Print Top-Level ExportedProgram
    # --------------------------------------------------------------------------
    print("\n[STEP 1] Top-Level Export: torch.export.export(..., strict=True)")
    ep = torch.export.export(model, args=example_args, strict=True)

    print("\n--- Top-Level ExportedProgram ---")
    print(ep)

    print(
        f"\n  --> Top-Level user_inputs_to_mutate: {ep.graph_signature.user_inputs_to_mutate}"
    )
    print("  --> Top-Level Output Specs:")
    for idx, spec in enumerate(ep.graph_signature.output_specs):
        print(f"        Output #{idx}: kind={spec.kind}, arg={spec.arg}")

    # --------------------------------------------------------------------------
    # 2. Decompose Graph & Print Decomposed ExportedProgram
    # --------------------------------------------------------------------------
    print(
        "\n[STEP 2] Decomposed ExportedProgram Verification (run_decompositions)"
    )
    decomposed_ep = ep.run_decompositions()

    print("\n--- Decomposed ExportedProgram ---")
    print(decomposed_ep)

    dec_sig = decomposed_ep.graph_signature
    print(
        f"\n  --> Decomposed user_inputs_to_mutate: {dec_sig.user_inputs_to_mutate}"
    )
    print("  --> Decomposed Output Specs:")
    for idx, spec in enumerate(dec_sig.output_specs):
        print(f"        Output #{idx}: kind={spec.kind}, arg={spec.arg}")

    mutated_input_names = list(dec_sig.user_inputs_to_mutate.values())
    assert "x" in mutated_input_names, f"FAIL: 'x' was not tracked in user_inputs_to_mutate for {name}!"
    print(
        f"  ✅ Signature Verification Passed: Decomposed signature maps 'x' to USER_INPUT_MUTATION."
    )

    # --------------------------------------------------------------------------
    # 3. AOTI Compilation, Runtime Execution & Scheduled Profiler
    # --------------------------------------------------------------------------
    print(
        "\n[STEP 3] AOTI Compilation, Value Mutation Check & Scheduled Profiling"
    )
    with tempfile.TemporaryDirectory() as tmpdir:
        pkg_path = f"{tmpdir}/model.pt2"
        compiled_pkg = torch._inductor.aoti_compile_and_package(
            decomposed_ep, package_path=pkg_path)
        runner = torch._inductor.aoti_load_package(compiled_pkg)

        # Verify single execution value update
        x_runtime = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
                                 device=device)
        x_original = x_runtime.clone()

        _ = runner(x_runtime)

        expected_x = x_original * 2
        assert torch.equal(
            x_runtime,
            expected_x), f"FAIL: Expected {expected_x}, got {x_runtime}"
        print(
            f"  ✅ Value Mutation Passed: Tensor mutated in-place to:\n{x_runtime}"
        )

        # --- Prepare Inputs and Profiler Schedule ---
        trace_filename = f"aoti_trace_{model_class.__name__}.json"

        activities = [torch.profiler.ProfilerActivity.CPU]
        if device == "cuda":
            activities.append(torch.profiler.ProfilerActivity.CUDA)

        wait_steps = 1
        warmup_steps = 2
        active_steps = 5
        total_steps = wait_steps + warmup_steps + active_steps

        prof_schedule = torch.profiler.schedule(wait=wait_steps,
                                                warmup=warmup_steps,
                                                active=active_steps,
                                                repeat=1)

        # Pre-allocate input tensor once on device (stateful mutation across steps)
        x_prof = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
                              device=device)

        with torch.profiler.profile(
                activities=activities,
                schedule=prof_schedule,
                record_shapes=True,
                with_stack=True,
        ) as prof:
            for _ in range(total_steps):
                runner(x_prof)
                prof.step()

        prof.export_chrome_trace(trace_filename)
        print(
            f"  ✅ Torch Profiler Trace Saved to: {os.path.abspath(trace_filename)}"
        )


def main() -> None:
    test_mutation_pipeline(
        UserInputMutationTriton,
        "Module WITH Custom Triton Kernel Mutation (triton_op/wrap_triton)")
    test_mutation_pipeline(UserInputMutationNoReturn,
                           "Module WITHOUT Return of Mutated Input")
    test_mutation_pipeline(UserInputMutationWithReturn,
                           "Module WITH Return of Mutated Input")


if __name__ == "__main__":
    main()

The input mutations can be confirmed by checking the user_inputs_to_mutate mapping in the graph signature of the decomposed ExportedProgram or just the decomposed ExportedProgram itself.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
$ python aoti_input_mutation_example.py

===========================================================================
 TESTING: Module WITH Custom Triton Kernel Mutation (triton_op/wrap_triton)
===========================================================================

[STEP 1] Top-Level Export: torch.export.export(..., strict=True)

--- Top-Level ExportedProgram ---
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "f32[3, 2]"):
            # File: /mnt/aoti_input_mutation_example.py:42 in forward, code: triton_mul2_mutation(x)
            triton_mul2_mutation = torch.ops.mylib.triton_mul2_mutation.default(x);  triton_mul2_mutation = None

            # File: /mnt/aoti_input_mutation_example.py:43 in forward, code: return x.cos()
            cos: "f32[3, 2]" = torch.ops.aten.cos.default(x);  x = None
            return (cos,)

Graph signature:
    # inputs
    x: USER_INPUT

    # outputs
    cos: USER_OUTPUT

Range constraints: {}


  --> Top-Level user_inputs_to_mutate: {}
  --> Top-Level Output Specs:
        Output #0: kind=OutputKind.USER_OUTPUT, arg=TensorArgument(name='cos')

[STEP 2] Decomposed ExportedProgram Verification (run_decompositions)
/usr/lib/python3.12/copyreg.py:99: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)

--- Decomposed ExportedProgram ---
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "f32[3, 2]"):
            # File: /mnt/aoti_input_mutation_example.py:42 in forward, code: triton_mul2_mutation(x)
            auto_functionalized_v2 = torch.ops.higher_order.auto_functionalized_v2(torch.ops.mylib.triton_mul2_mutation.default, _x_base_index = 0, _all_bases = [x]);  x = None
            getitem_1: "f32[3, 2]" = auto_functionalized_v2[1];  auto_functionalized_v2 = None

            # File: /mnt/aoti_input_mutation_example.py:43 in forward, code: return x.cos()
            cos: "f32[3, 2]" = torch.ops.aten.cos.default(getitem_1)
            return (getitem_1, cos)

Graph signature:
    # inputs
    x: USER_INPUT

    # outputs
    getitem_1: USER_INPUT_MUTATION target='x'
    cos: USER_OUTPUT

Range constraints: {}


  --> Decomposed user_inputs_to_mutate: {'getitem_1': 'x'}
  --> Decomposed Output Specs:
        Output #0: kind=OutputKind.USER_INPUT_MUTATION, arg=TensorArgument(name='getitem_1')
        Output #1: kind=OutputKind.USER_OUTPUT, arg=TensorArgument(name='cos')
  ✅ Signature Verification Passed: Decomposed signature maps 'x' to USER_INPUT_MUTATION.

[STEP 3] AOTI Compilation, Value Mutation Check & Scheduled Profiling
/usr/lib/python3.12/copyreg.py:99: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)
  ✅ Value Mutation Passed: Tensor mutated in-place to:
tensor([[ 2.,  4.],
        [ 6.,  8.],
        [10., 12.]], device='cuda:0')
USDT:2026-09-01 05:58:14 230:230 SyncActivityProfilerHandler.cpp:52] profiler_start
USDT:2026-09-01 05:58:14 230:230 SyncActivityProfilerHandler.cpp:59] profiler_stop
  ✅ Torch Profiler Trace Saved to: /mnt/aoti_trace_UserInputMutationTriton.json

===========================================================================
 TESTING: Module WITHOUT Return of Mutated Input
===========================================================================

[STEP 1] Top-Level Export: torch.export.export(..., strict=True)

--- Top-Level ExportedProgram ---
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "f32[3, 2]"):
            # File: /mnt/aoti_input_mutation_example.py:50 in forward, code: x.mul_(2)
            mul_: "f32[3, 2]" = torch.ops.aten.mul_.Tensor(x, 2);  x = None

            # File: /mnt/aoti_input_mutation_example.py:51 in forward, code: return x.cos()
            cos: "f32[3, 2]" = torch.ops.aten.cos.default(mul_);  mul_ = None
            return (cos,)

Graph signature:
    # inputs
    x: USER_INPUT

    # outputs
    cos: USER_OUTPUT

Range constraints: {}


  --> Top-Level user_inputs_to_mutate: {}
  --> Top-Level Output Specs:
        Output #0: kind=OutputKind.USER_OUTPUT, arg=TensorArgument(name='cos')

[STEP 2] Decomposed ExportedProgram Verification (run_decompositions)
/usr/lib/python3.12/copyreg.py:99: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)

--- Decomposed ExportedProgram ---
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "f32[3, 2]"):
            # File: /mnt/aoti_input_mutation_example.py:50 in forward, code: x.mul_(2)
            mul: "f32[3, 2]" = torch.ops.aten.mul.Tensor(x, 2);  x = None

            # File: /mnt/aoti_input_mutation_example.py:51 in forward, code: return x.cos()
            cos: "f32[3, 2]" = torch.ops.aten.cos.default(mul)
            return (mul, cos)

Graph signature:
    # inputs
    x: USER_INPUT

    # outputs
    mul: USER_INPUT_MUTATION target='x'
    cos: USER_OUTPUT

Range constraints: {}


  --> Decomposed user_inputs_to_mutate: {'mul': 'x'}
  --> Decomposed Output Specs:
        Output #0: kind=OutputKind.USER_INPUT_MUTATION, arg=TensorArgument(name='mul')
        Output #1: kind=OutputKind.USER_OUTPUT, arg=TensorArgument(name='cos')
  ✅ Signature Verification Passed: Decomposed signature maps 'x' to USER_INPUT_MUTATION.

[STEP 3] AOTI Compilation, Value Mutation Check & Scheduled Profiling
/usr/lib/python3.12/copyreg.py:99: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)
  ✅ Value Mutation Passed: Tensor mutated in-place to:
tensor([[ 2.,  4.],
        [ 6.,  8.],
        [10., 12.]], device='cuda:0')
USDT:2026-09-01 05:58:19 230:230 SyncActivityProfilerHandler.cpp:52] profiler_start
USDT:2026-09-01 05:58:19 230:230 SyncActivityProfilerHandler.cpp:59] profiler_stop
  ✅ Torch Profiler Trace Saved to: /mnt/aoti_trace_UserInputMutationNoReturn.json

===========================================================================
 TESTING: Module WITH Return of Mutated Input
===========================================================================

[STEP 1] Top-Level Export: torch.export.export(..., strict=True)

--- Top-Level ExportedProgram ---
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "f32[3, 2]"):
            # File: /mnt/aoti_input_mutation_example.py:58 in forward, code: x.mul_(2)
            mul_: "f32[3, 2]" = torch.ops.aten.mul_.Tensor(x, 2);  x = None

            # File: /mnt/aoti_input_mutation_example.py:59 in forward, code: return x, x.cos()
            cos: "f32[3, 2]" = torch.ops.aten.cos.default(mul_)
            return (mul_, cos)

Graph signature:
    # inputs
    x: USER_INPUT

    # outputs
    mul_: USER_OUTPUT
    cos: USER_OUTPUT

Range constraints: {}


  --> Top-Level user_inputs_to_mutate: {}
  --> Top-Level Output Specs:
        Output #0: kind=OutputKind.USER_OUTPUT, arg=TensorArgument(name='mul_')
        Output #1: kind=OutputKind.USER_OUTPUT, arg=TensorArgument(name='cos')

[STEP 2] Decomposed ExportedProgram Verification (run_decompositions)
/usr/lib/python3.12/copyreg.py:99: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)

--- Decomposed ExportedProgram ---
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "f32[3, 2]"):
            # File: /mnt/aoti_input_mutation_example.py:58 in forward, code: x.mul_(2)
            mul: "f32[3, 2]" = torch.ops.aten.mul.Tensor(x, 2);  x = None

            # File: /mnt/aoti_input_mutation_example.py:59 in forward, code: return x, x.cos()
            cos: "f32[3, 2]" = torch.ops.aten.cos.default(mul)
            return (mul, mul, cos)

Graph signature:
    # inputs
    x: USER_INPUT

    # outputs
    mul: USER_INPUT_MUTATION target='x'
    mul: USER_OUTPUT
    cos: USER_OUTPUT

Range constraints: {}


  --> Decomposed user_inputs_to_mutate: {'mul': 'x'}
  --> Decomposed Output Specs:
        Output #0: kind=OutputKind.USER_INPUT_MUTATION, arg=TensorArgument(name='mul')
        Output #1: kind=OutputKind.USER_OUTPUT, arg=TensorArgument(name='mul')
        Output #2: kind=OutputKind.USER_OUTPUT, arg=TensorArgument(name='cos')
  ✅ Signature Verification Passed: Decomposed signature maps 'x' to USER_INPUT_MUTATION.

[STEP 3] AOTI Compilation, Value Mutation Check & Scheduled Profiling
/usr/lib/python3.12/copyreg.py:99: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)
  ✅ Value Mutation Passed: Tensor mutated in-place to:
tensor([[ 2.,  4.],
        [ 6.,  8.],
        [10., 12.]], device='cuda:0')
USDT:2026-09-01 05:58:23 230:230 SyncActivityProfilerHandler.cpp:52] profiler_start
USDT:2026-09-01 05:58:23 230:230 SyncActivityProfilerHandler.cpp:59] profiler_stop
  ✅ Torch Profiler Trace Saved to: /mnt/aoti_trace_UserInputMutationWithReturn.json

Caveats

In many cases, the inputs being mutated are caches. A natural implementation would just create a PyTorch model that has internal buffers registered via self.register_buffer(...) and mutate those buffers in-place. The AOTInductor engine generated from such a model will have thread-safety issues, which is not immediately obvious, if multiple threads are running the same engine concurrently, because the buffers are shared across threads. The input mutation approach, however, is thread-safe, because each thread has its own input tensor to mutate.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import os
import tempfile
import torch


class EagerMutationModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # A persistent buffer that tracks execution state
        self.register_buffer("counter", torch.zeros(1))

    def forward(self, x):
        # IN-PLACE MUTATION: State update on a registered buffer
        self.counter.add_(1)
        return x * self.counter


def test_buffer_mutation_pipeline() -> None:
    print(f"\n{'='*75}")
    print(" TESTING BUFFER MUTATION: EagerMutationModel with AOTI & Profiler")
    print(f"{'='*75}")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = EagerMutationModel().to(device).eval()
    example_args = (torch.randn(3, 2, device=device),)

    # --------------------------------------------------------------------------
    # 1. Export with strict=True & Print Top-Level ExportedProgram
    # --------------------------------------------------------------------------
    print("\n[STEP 1] Top-Level Export: torch.export.export(..., strict=True)")
    ep = torch.export.export(model, args=example_args, strict=True)

    print("\n--- Top-Level ExportedProgram ---")
    print(ep)

    print(f"\n  --> Top-Level buffers_to_mutate: {ep.graph_signature.buffers_to_mutate}")
    print("  --> Top-Level Output Specs:")
    for idx, spec in enumerate(ep.graph_signature.output_specs):
        print(f"        Output #{idx}: kind={spec.kind}, arg={spec.arg}")

    # --------------------------------------------------------------------------
    # 2. Decompose Graph & Print Decomposed ExportedProgram
    # --------------------------------------------------------------------------
    print("\n[STEP 2] Decomposed ExportedProgram Verification (run_decompositions)")
    decomposed_ep = ep.run_decompositions()

    print("\n--- Decomposed ExportedProgram ---")
    print(decomposed_ep)

    dec_sig = decomposed_ep.graph_signature
    print(f"\n  --> Decomposed buffers_to_mutate: {dec_sig.buffers_to_mutate}")
    print("  --> Decomposed Output Specs:")
    for idx, spec in enumerate(dec_sig.output_specs):
        print(f"        Output #{idx}: kind={spec.kind}, arg={spec.arg}")

    mutated_buffer_names = list(dec_sig.buffers_to_mutate.values())
    assert "counter" in mutated_buffer_names, "FAIL: 'counter' was not tracked in buffers_to_mutate!"
    print(f"\n  ✅ Signature Verification Passed: Decomposed signature maps 'counter' to BUFFER_MUTATION.")

    # --------------------------------------------------------------------------
    # 3. AOTI Compilation, Runtime Execution & Scheduled Profiler
    # --------------------------------------------------------------------------
    print("\n[STEP 3] AOTI Compilation, Buffer Mutation Check & Scheduled Profiling")
    with tempfile.TemporaryDirectory() as tmpdir:
        pkg_path = f"{tmpdir}/model.pt2"
        compiled_pkg = torch._inductor.aoti_compile_and_package(
            decomposed_ep, package_path=pkg_path
        )
        runner = torch._inductor.aoti_load_package(compiled_pkg)

        # Verify initial execution state update
        x_runtime = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], device=device)
        _ = runner(x_runtime)
        print(f"  ✅ AOTI Execution Passed: Buffer state updated successfully.")

        # --- Prepare Inputs and Profiler Schedule ---
        trace_filename = "aoti_trace_EagerMutationModel.json"

        activities = [torch.profiler.ProfilerActivity.CPU]
        if device == "cuda":
            activities.append(torch.profiler.ProfilerActivity.CUDA)

        wait_steps = 1
        warmup_steps = 2
        active_steps = 5
        total_steps = wait_steps + warmup_steps + active_steps

        prof_schedule = torch.profiler.schedule(
            wait=wait_steps,
            warmup=warmup_steps,
            active=active_steps,
            repeat=1
        )

        # Pre-allocate input tensor once on device before profiling loop
        x_prof = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], device=device)

        with torch.profiler.profile(
            activities=activities,
            schedule=prof_schedule,
            record_shapes=True,
            with_stack=True,
        ) as prof:
            for _ in range(total_steps):
                runner(x_prof)
                prof.step()

        prof.export_chrome_trace(trace_filename)
        print(f"  ✅ Torch Profiler Trace Saved to: {os.path.abspath(trace_filename)}")


if __name__ == "__main__":
    test_buffer_mutation_pipeline()

In the decomposed ExportedProgram, we could clearly see that the buffer counter is tracked in buffers_to_mutate, and the output spec for the mutated buffer is of kind BUFFER_MUTATION. Consequently, we should avoid mutating registered buffers in the PyTorch model.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
$ python buffer_mutation.py

===========================================================================
 TESTING BUFFER MUTATION: EagerMutationModel with AOTI & Profiler
===========================================================================

[STEP 1] Top-Level Export: torch.export.export(..., strict=True)

--- Top-Level ExportedProgram ---
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, b_counter: "f32[1]", x: "f32[3, 2]"):
            # File: /mnt/i.py:14 in forward, code: self.counter.add_(1)
            add_: "f32[1]" = torch.ops.aten.add_.Tensor(b_counter, 1);  b_counter = None

            # File: /mnt/i.py:15 in forward, code: return x * self.counter
            mul: "f32[3, 2]" = torch.ops.aten.mul.Tensor(x, add_);  x = add_ = None
            return (mul,)

Graph signature:
    # inputs
    b_counter: BUFFER target='counter' persistent=True
    x: USER_INPUT

    # outputs
    mul: USER_OUTPUT

Range constraints: {}


  --> Top-Level buffers_to_mutate: {}
  --> Top-Level Output Specs:
        Output #0: kind=OutputKind.USER_OUTPUT, arg=TensorArgument(name='mul')

[STEP 2] Decomposed ExportedProgram Verification (run_decompositions)
/usr/lib/python3.12/copyreg.py:99: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)

--- Decomposed ExportedProgram ---
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, b_counter: "f32[1]", x: "f32[3, 2]"):
            # File: /mnt/i.py:14 in forward, code: self.counter.add_(1)
            add: "f32[1]" = torch.ops.aten.add.Tensor(b_counter, 1);  b_counter = None

            # File: /mnt/i.py:15 in forward, code: return x * self.counter
            mul: "f32[3, 2]" = torch.ops.aten.mul.Tensor(x, add);  x = None
            return (add, mul)

Graph signature:
    # inputs
    b_counter: BUFFER target='counter' persistent=True
    x: USER_INPUT

    # outputs
    add: BUFFER_MUTATION target='counter'
    mul: USER_OUTPUT

Range constraints: {}


  --> Decomposed buffers_to_mutate: {'add': 'counter'}
  --> Decomposed Output Specs:
        Output #0: kind=OutputKind.BUFFER_MUTATION, arg=TensorArgument(name='add')
        Output #1: kind=OutputKind.USER_OUTPUT, arg=TensorArgument(name='mul')

  ✅ Signature Verification Passed: Decomposed signature maps 'counter' to BUFFER_MUTATION.

[STEP 3] AOTI Compilation, Buffer Mutation Check & Scheduled Profiling
/usr/lib/python3.12/copyreg.py:99: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)
  ✅ AOTI Execution Passed: Buffer state updated successfully.
USDT:2026-08-29 04:06:00 4197:4197 SyncActivityProfilerHandler.cpp:52] profiler_start
USDT:2026-08-29 04:06:00 4197:4197 SyncActivityProfilerHandler.cpp:59] profiler_stop
  ✅ Torch Profiler Trace Saved to: /mnt/aoti_trace_EagerMutationModel.json

Conclusions

TensorRT allows input mutations, so does AOTInductor.

References

AOTInductor Input Mutation

https://leimao.github.io/blog/2026/2026-09-01-AOTInductor-Input-Mutation/

Author

Lei Mao

Posted on

09-01-2026

Updated on

09-01-2026

Licensed under

 

Like this article? Support the author with

Comments

  Follow   Sponsor

Advertisement


Catalogue

×