Allows boolean attn_mask.

xinyazhang · xinyazhang · commit 1489d358d1aa · 2025-01-06T18:50:30.000Z
diff --git a/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp b/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp
@@ -708,8 +708,10 @@ bool can_use_mem_efficient_attention(sdp_params const& params, bool debug) {
 
 #ifdef USE_ROCM
   if (params.attn_mask.has_value()) {
-    if (params.attn_mask.value().dtype() != params.query.dtype()) {
-      TORCH_WARN("Efficient attention on ROCM requires attn_mask has the same datatype as of q,k,v");
+    const auto q_dtype = params.query.dtype();
+    const auto bias_dtype = params.attn_mask.value().dtype();
+    if (bias_dtype != at::kBool && bias_dtype != q_dtype) {
+      TORCH_WARN("Efficient attention on ROCM requires attn_mask be boolean, or has the same datatype as of q,k,v");
       return false;
     }
   }

Original file line number	Diff line number	Diff line change
`@@ -708,8 +708,10 @@ bool can_use_mem_efficient_attention(sdp_params const& params, bool debug) {`
`708`	`708`
`709`	`709`	`#ifdef USE_ROCM`
`710`	`710`	`if (params.attn_mask.has_value()) {`
`711`		`- if (params.attn_mask.value().dtype() != params.query.dtype()) {`
`712`		`- TORCH_WARN("Efficient attention on ROCM requires attn_mask has the same datatype as of q,k,v");`
	`711`	`+ const auto q_dtype = params.query.dtype();`
	`712`	`+ const auto bias_dtype = params.attn_mask.value().dtype();`
	`713`	`+ if (bias_dtype != at::kBool && bias_dtype != q_dtype) {`
	`714`	`+ TORCH_WARN("Efficient attention on ROCM requires attn_mask be boolean, or has the same datatype as of q,k,v");`
`713`	`715`	`return false;`
`714`	`716`	`}`
`715`	`717`	`}`