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.
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.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124
|
import osimport tempfileimport torchclass 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(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.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153
|
$ python aoti_input_mutation_example.py=========================================================================== 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/g.py:8 in forward, code: x.mul_(2) mul_: "f32[3, 2]" = torch.ops.aten.mul_.Tensor(x, 2); x = None # File: /mnt/g.py:9 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_OUTPUTRange 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/g.py:8 in forward, code: x.mul_(2) mul: "f32[3, 2]" = torch.ops.aten.mul.Tensor(x, 2); x = None # File: /mnt/g.py:9 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_OUTPUTRange 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-08-29 03:44:32 3839:3839 SyncActivityProfilerHandler.cpp:52] profiler_startUSDT:2026-08-29 03:44:32 3839:3839 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/g.py:15 in forward, code: x.mul_(2) mul_: "f32[3, 2]" = torch.ops.aten.mul_.Tensor(x, 2); x = None # File: /mnt/g.py:16 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_OUTPUTRange 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/g.py:15 in forward, code: x.mul_(2) mul: "f32[3, 2]" = torch.ops.aten.mul.Tensor(x, 2); x = None # File: /mnt/g.py:16 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_OUTPUTRange 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-08-29 03:44:37 3839:3839 SyncActivityProfilerHandler.cpp:52] profiler_startUSDT:2026-08-29 03:44:37 3839:3839 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.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113
|
import osimport tempfileimport torchclass 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.counterdef 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.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
|
$ 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_OUTPUTRange 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_OUTPUTRange 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_startUSDT: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/
source & further reading
leimao.github.io — original article
AOTInductor External Weight Storage and Weight Streaming Update
CUDA Graph In The Context of Multi-Stream Execution
PyTorch Asynchronous Assert