open mask

example/32_batched_gemm_scale_softmax_gemm/batched_multihead_attention_backward_fp16.cpp

@@ -24,6 +24,7 @@ Kernel outputs:
 */
 
 #define PRINT_HOST 0
+#define USING_MASK 1
 
 #include <iostream>
 #include <numeric>
@@ -69,8 +70,13 @@ static constexpr ck::index_t NumDimK = 1;
 static constexpr ck::index_t NumDimO = 1;
 
 static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::MNKOPadding;
+#if USING_MASK
 static constexpr auto MaskingSpec =
-    ck::tensor_operation::device::MaskingSpecialization::MaskDisabled;
+    ck::tensor_operation::device::MaskingSpecialization::MaskOutUpperTriangle;
+#else
+static constexpr auto MaskingSpec =
+    ck::tensor_operation::device::MaskingSpecialization::MaskOutUpperTriangle;
+#endif
 
 static constexpr auto TensorSpecQ = ck::tensor_operation::device::TensorSpecialization::Default;
 static constexpr auto TensorSpecK = ck::tensor_operation::device::TensorSpecialization::Default;
@@ -203,7 +209,8 @@ void run_attention_fwd_host(const TensorQ& q_g_m_k,
     ref_gemm0_invoker.Run(ref_gemm0_argument);
 
     // masking
-#if 0
+#if USING_MASK
+    auto N          = s_g_m_n.GetLengths()[1];
     const auto mask = DeviceGemmInstance::C0MatrixMask(N);
     s_g_m_n.ForEach([&](auto& self, auto idx) {
         if(mask.IsMaskedElement(idx[1], idx[2]))