Flops measurement for PT fMHA operators (facebookresearch#1127)

bertmaher · Jun 7, 2024 · c4e8cc0 · c4e8cc0
1 parent 7d111fd
commit c4e8cc0
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/xformers/profiler/profile_analyzer.py b/xformers/profiler/profile_analyzer.py
@@ -64,10 +64,12 @@ def _replace_if_needed(
             (torch.ops.aten, "scaled_dot_product_attention", False),
             (torch.ops.xformers_flash, "flash_fwd", False),
             (torch.ops.xformers, "efficient_attention_forward_cutlass", False),
+            (torch.ops.aten, "_efficient_attention_forward", False),
             (torch.ops.aten, "_scaled_dot_product_flash_attention_backward", True),
             (torch.ops.aten, "_scaled_dot_product_efficient_attention_backward", True),
             (torch.ops.xformers_flash, "flash_bwd", True),
             (torch.ops.xformers, "efficient_attention_backward_cutlass", True),
+            (torch.ops.aten, "_efficient_attention_backward", True),
         ]
         if hasattr(lib, op)
     }