Disable donated buffer when benchmarking

pytorch-labs · Dec 2, 2024 · afa8e31 · afa8e31
1 parent c509d84
commit afa8e31
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tritonbench/operators/layer_norm/operator.py b/tritonbench/operators/layer_norm/operator.py
@@ -34,6 +34,11 @@ def torch_layer_norm(self, *args):
 
     @register_benchmark()
     def torch_compile_layer_norm(self, *args):
+        # We need to run backward multiple times for proper benchmarking
+        # so donated buffer have to be disabled
+        if self.mode == Mode.BWD or self.mode == Mode.FWD_BWD:
+            import torch._functorch.config
+            torch._functorch.config.donated_buffer = False
         @torch.compile
         def inner(*args):
             return F.layer_norm(*args)