[inductor] estimate peak memory in codegen only when buffer reuse (#162300)

ruisizhang123 · pytorchmergebot · commit 291cd11f2d5d · 2025-09-06T01:30:38.000Z
As titled, this PR ensures peak memory is estimated only when buffer reuse is enabled. Without this config, some nodes' successor nodes are eliminated from memory estimation after inductor bucketing, which can cause errors. The original codegen peak memory estimation code is from this PR: #159530 Pull Request resolved: #162300 Approved by: https://github.com/eellison, https://github.com/v0i0
diff --git a/torch/_inductor/codegen/wrapper.py b/torch/_inductor/codegen/wrapper.py
@@ -1736,7 +1736,8 @@ def run_wrapper_ir_passes(self, is_inference: bool):
         if is_inference and config.memory_planning:
             self.memory_plan()
         else:
-            self.estimate_peak = EfficientPeakEstimate()
+            if config.allow_buffer_reuse:
+                self.estimate_peak = EfficientPeakEstimate()
             self.memory_plan_reuse()
 
     def codegen_input_symbol_assignment(
