memory_efficient_attention() nondeterministic warning (#635)

* add warning of non-deterministic behavior
for efficient_attention_forward_cutlass

* add non-deterministic note

* add alertNotDeterministic to
mem_efficient_attention_backward_cutlass

* Update document, thanks @danthe3rd

xformers/csrc/attention/cuda/fmha/attention_backward_generic.cu

@@ -4,6 +4,7 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <torch/library.h>
+#include <ATen/Context.h>
 
 #include "kernel_backward.h"
 
@@ -67,6 +68,8 @@ mem_efficient_attention_backward_cutlass(
       false,
       "MemoryEfficient build has been disabled at build time with -DXFORMERS_MEM_EFF_ATTENTION_DISABLE_BACKWARD");
 #else
+  at::globalContext().alertNotDeterministic("mem_efficient_attention_backward_cutlass");
+
   // ndim
   TORCH_CHECK(query.dim() == grad_out_.dim());
   TORCH_CHECK(query.dim() == key.dim());

xformers/csrc/attention/cuda/fmha/attention_forward_generic.cu

@@ -3,6 +3,7 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <torch/library.h>
+#include <ATen/Context.h>
 
 #include "kernel_forward.h"
 
@@ -144,6 +145,8 @@ std::tuple<at::Tensor, at::Tensor> efficient_attention_forward_cutlass(
       false,
       "MemoryEfficient build has been disabled at build time with -DXFORMERS_MEM_EFF_ATTENTION_DISABLE_FORWARD");
 #else
+  at::globalContext().alertNotDeterministic("efficient_attention_forward_cutlass");
+
   TORCH_CHECK(query.dim() == 4);
   TORCH_CHECK(key.dim() == 4);
   TORCH_CHECK(value.dim() == 4);

xformers/ops/fmha/__init__.py

@@ -172,6 +172,10 @@ def memory_efficient_attention(
 
         NVIDIA GPUs with compute capability above 6.0 (P100+), datatype ``f16``, ``bf16`` and ``f32``.
 
+    :Note:
+
+        This operator may be nondeterministic.
+
     Raises:
         NotImplementedError: if there is no operator available to compute the MHA