[CPU SDP] Remove mem efficient attn checks in CPU (#112375)

rohan-varma · pytorchmergebot · commit 92cc52ab0e48 · 2023-10-30T16:43:20.000Z
It doesn't seem like memory efficient attention can be used on CPU, as we don't check for it when iterating backends in `select_sdp_backend_cpp`. So removing some of the logic around mem efficient attention selection. Created from CodeHub with https://fburl.com/edit-in-codehub Differential Revision: [D50775562](https://our.internmc.facebook.com/intern/diff/D50775562/) **NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D50775562/)! Pull Request resolved: #112375 Approved by: https://github.com/drisspg
diff --git a/aten/src/ATen/native/transformers/sdp_utils_cpp.cpp b/aten/src/ATen/native/transformers/sdp_utils_cpp.cpp
@@ -61,11 +61,9 @@ bool use_flash_attention_cpp(sdp_params const& params, bool debug) {
 SDPBackend select_sdp_backend_cpp(sdp_params const& kernel_params) {
   // This function defines the priority order of the different sdp backends
   // 1. Flash Attention
-  // 2. Mem Efficient Attention
-  // 3. Math fallback
+  // 2. Math fallback
   auto& ctx = at::globalContext();
-  if (!ctx.userEnabledMathSDP() && !ctx.userEnabledFlashSDP() &&
-      !ctx.userEnabledMemEfficientSDP()) {
+  if (!ctx.userEnabledMathSDP() && !ctx.userEnabledFlashSDP()) {
     return SDPBackend::error;
   }
   // Get ideal kernel ordering
@@ -91,7 +89,7 @@ SDPBackend select_sdp_backend_cpp(sdp_params const& kernel_params) {
     }
   }
   // If we have gotten to this point then two things have happened:
-  // 1. use_flash_attention or use_mem_efficient did not satisfy the
+  // 1. use_flash_attention did not satisfy the
   // constraints to be ran
   // 2. The user has explicitly disabled the math kernel
   // We then re-run the kernel checks with debug enabled to print out the
