Flops measurement for PT fMHA operators (fairinternal/xformers#1127)

__original_commit__ = fairinternal/xformers@c4e8cc0
facebookresearch · Jun 7, 2024 · cda3a4e · cda3a4e
1 parent 58d7a65
commit cda3a4e
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)
     }