[inductor][invoke_subgraph] Free the buffers before the subgraph call (#152494)

anijain2305 · pytorchmergebot · commit 5236a8506c4f · 2025-05-01T02:04:10.000Z
Before ![image](https://github.com/user-attachments/assets/62b24c14-69e6-40fb-94e3-223930132ef6) After ![image](https://github.com/user-attachments/assets/9f340d4e-80a9-45aa-9400-626fff5b5ecd) tlparse - https://manifold.edge.x2p.facebook.net/v0/read/tree/logs/.tmph5dwWt/index.html?bucketName=tlparse_reports&apiKey=tlparse_reports-key&withPayload=1&timeoutMsec=10000 Pull Request resolved: #152494 Approved by: https://github.com/Skylion007, https://github.com/eellison ghstack dependencies: #152357, #152384, #152383, #152490
diff --git a/torch/_inductor/codegen/wrapper.py b/torch/_inductor/codegen/wrapper.py
@@ -2909,6 +2909,10 @@ def codegen_subgraph_call(self, subgraph, outer_inputs, outer_buffer_name):
 
         self.writeline(f"{subgraph.graph.name}_args = [{outer_input_names}]")
 
+        # Since the buffers are already put into the args list, we can free the
+        # buffers here.
+        V.graph.scheduler.free_buffers()
+
         # Call the subgraph launcher function
         self.writeline(
             f"{outer_buffer_name} = {subgraph.graph.name}({subgraph.graph.name}_args)"
