Enable get dropout_mask

src/ATen/native/transformers/Attention.cpp

@@ -3,6 +3,8 @@
 #include <ATen/native/nested/NestedTensorUtils.h>
 #include <ATen/native/transformers/attention.h>
 #include <ATen/native/transformers/sdp_utils_cpp.h>
+#include <ATen/native/xpu/sycl/DropoutKernels.h>
+#include <ATen/ops/_fill_mem_eff_dropout_mask_native.h>
 
 #ifndef AT_PER_OPERATOR_HEADERS
 #include <ATen/Functions.h>
@@ -294,5 +296,19 @@ std::tuple<Tensor, Tensor> native_multi_head_attention_xpu(
   return std::make_tuple(std::move(proj), std::move(qkt));
 }
 
+/**
+ * get the mask for dropout. only used for testing, not much
+ * attention is paid to performance
+ */
+at::Tensor& _fill_mem_eff_dropout_mask_(
+    Tensor& self,
+    double dropout_p,
+    const int64_t seed,
+    const int64_t offset) {
+  auto mask = std::get<1>(xpu::dropout_kernel(self, dropout_p, true));
+  self.copy_(mask);
+  return self;
+}
+
 } // namespace native
 } // namespace at

yaml/native/native_functions.yaml

@@ -7922,6 +7922,12 @@
     SparseCsrXPU: angle_sparse_csr_out
   tags: pointwise
 
+- func: _fill_mem_eff_dropout_mask_(Tensor(a!) self, float dropout_p, int seed, int offset) -> Tensor(a!)
+  variants: function
+  dispatch:
+    XPU: _fill_mem_eff_dropout_mask_
+  tags: nondeterministic_seeded
+
 - func: special_airy_ai(Tensor x) -> Tensor
   python_module: special
   structured_delegate: special_airy_ai.out