AOTInductor Input Mutation

AOTInductor, a PyTorch compiler backend, supports in-place input mutations in optimized inference engines, contrary to the assumption that it relies solely on functionalization. The key is to explicitly use in-place operations like x.mul_(2) in the model; after decomposition, the graph signature tracks the mutation, and AOTInductor generates an engine that performs the in-place mutation. This enables efficient updates to small portions of large tensors without allocating new tensors.

AOTInductor Input Mutation 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 | python 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 | bash $ 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 | python 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 | bash $ 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/ https://leimao.github.io/blog/2026/2026-09-01-AOTInductor-Input-Mutation/