make causal mask/bias both optional

openxla · Jan 12, 2024 · 73928b4 · 73928b4
1 parent 9af8953
commit 73928b4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/xla/stream_executor/cuda/cuda_dnn.cc b/xla/stream_executor/cuda/cuda_dnn.cc
@@ -9427,7 +9427,8 @@ CudnnSupport::FusedMHABackwardRunnerFromDesc(
         CudnnfMHAUid::d_Q_accum_ID, CudnnfMHAUid::O_ID};
     if (bias_descriptor != std::nullopt) {
       uids.push_back(CudnnfMHAUid::BIAS_ID);
-    } else {
+    }
+    if (is_causal_mask) {
       // is causal mask
       // negative infinity
       double negative_infinity_value = -std::numeric_limits<float>::infinity();