Merge pull request #811 from ROCmSoftwarePlatform/attn-train-develop-qloop-mask

Add another mask(upper tringle from bottom right) to flash attetion

profiler/include/profiler/profile_batched_gemm_softmax_gemm_permute_impl.hpp

@@ -219,7 +219,8 @@ bool profile_batched_gemm_softmax_gemm_permute_impl(bool do_verification,
 
         // mask out upper triangle
         acc0_g_m_n.ForEach([&](auto& self, auto idx) {
-            if(MaskingSpec == MaskingSpecialization::MaskOutUpperTriangle && idx[1] < idx[2])
+            if(MaskingSpec == MaskingSpecialization::MaskUpperTriangleFromTopLeft &&
+               idx[1] < idx[2])
                 self(idx) = -ck::NumericLimits<float>::Infinity();
         });
 

test/batched_gemm_softmax_gemm_permute/test_batched_gemm_bias_softmax_gemm_permute_bf16.cpp

@@ -15,13 +15,14 @@ using I2_t = ck::Number<2>;
 
 using MaskDisabled_t =
     ck::integral_constant<MaskingSpecialization, MaskingSpecialization::MaskDisabled>;
-using MaskOutUpperTriangle_t =
-    ck::integral_constant<MaskingSpecialization, MaskingSpecialization::MaskOutUpperTriangle>;
+using MaskUpperTriangleFromTopLeft_t =
+    ck::integral_constant<MaskingSpecialization,
+                          MaskingSpecialization::MaskUpperTriangleFromTopLeft>;
 
 // clang-format off
 using KernelTypes = ::testing::Types<
     std::tuple<I2_t, I1_t, I1_t, I1_t, I1_t, BF16, BF16, BF16, BF16, ck::Tuple<BF16>, ck::Tuple<>, MaskDisabled_t>,
-    std::tuple<I2_t, I1_t, I1_t, I1_t, I1_t, BF16, BF16, BF16, BF16, ck::Tuple<BF16>, ck::Tuple<>, MaskOutUpperTriangle_t>
+    std::tuple<I2_t, I1_t, I1_t, I1_t, I1_t, BF16, BF16, BF16, BF16, ck::Tuple<BF16>, ck::Tuple<>, MaskUpperTriangleFromTopLeft_t>
     >;
 // clang-format on
 

test/batched_gemm_softmax_gemm_permute/test_batched_gemm_bias_softmax_gemm_permute_fp16.cpp

@@ -15,13 +15,14 @@ using I2_t = ck::Number<2>;
 
 using MaskDisabled_t =
     ck::integral_constant<MaskingSpecialization, MaskingSpecialization::MaskDisabled>;
-using MaskOutUpperTriangle_t =
-    ck::integral_constant<MaskingSpecialization, MaskingSpecialization::MaskOutUpperTriangle>;
+using MaskUpperTriangleFromTopLeft_t =
+    ck::integral_constant<MaskingSpecialization,
+                          MaskingSpecialization::MaskUpperTriangleFromTopLeft>;
 
 // clang-format off
 using KernelTypes = ::testing::Types<
     std::tuple<I2_t, I1_t, I1_t, I1_t, I1_t, F16, F16, F16, F16, ck::Tuple<F16>, ck::Tuple<>, MaskDisabled_t>,
-    std::tuple<I2_t, I1_t, I1_t, I1_t, I1_t, F16, F16, F16, F16, ck::Tuple<F16>, ck::Tuple<>, MaskOutUpperTriangle_t>
+    std::tuple<I2_t, I1_t, I1_t, I1_t, I1_t, F16, F16, F16, F16, ck::Tuple<F16>, ck::Tuple<>, MaskUpperTriangleFromTopLeft_t>
     >;
 // clang-format on
 

test/batched_gemm_softmax_gemm_permute/test_batched_gemm_bias_softmax_gemm_permute_util.hpp

@@ -174,7 +174,7 @@ struct DeviceInstanceWrapper_G2M1N1K1O1_TNTT_FP16_M128_N128_K32_O128
             2,              // CShuffleNXdlPerWavePerShuffle
             S<1, 32, 1, 8>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock
             8,              // CShuffleBlockTransferScalarPerVector_NPerBlock
-            MaskingSpecialization::MaskOutUpperTriangle>; // MaskOutUpperTriangle
+            MaskingSpecialization::MaskUpperTriangleFromTopLeft>; // MaskUpperTriangleFromTopLeft
 
     bool IsSupported(int M, int N, int K, int O)
     {
@@ -321,7 +321,7 @@ struct DeviceInstanceWrapper_G2M1N1K1O1_TNTT_BF16_M128_N128_K32_O128
             2,              // CShuffleNXdlPerWavePerShuffle
             S<1, 32, 1, 8>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock
             8,              // CShuffleBlockTransferScalarPerVector_NPerBlock
-            MaskingSpecialization::MaskOutUpperTriangle>; // MaskOutUpperTriangle
+            MaskingSpecialization::MaskUpperTriangleFromTopLeft>; // MaskUpperTriangleFromTopLeft
 
     bool IsSupported(int M, int N, int K, int O)
     {

test/batched_gemm_softmax_gemm_permute/test_batched_gemm_softmax_gemm_permute_bf16.cpp

@@ -15,13 +15,14 @@ using I2_t = ck::Number<2>;
 
 using MaskDisabled_t =
     ck::integral_constant<MaskingSpecialization, MaskingSpecialization::MaskDisabled>;
-using MaskOutUpperTriangle_t =
-    ck::integral_constant<MaskingSpecialization, MaskingSpecialization::MaskOutUpperTriangle>;
+using MaskUpperTriangleFromTopLeft_t =
+    ck::integral_constant<MaskingSpecialization,
+                          MaskingSpecialization::MaskUpperTriangleFromTopLeft>;
 
 // clang-format off
 using KernelTypes = ::testing::Types<
     std::tuple<I2_t, I1_t, I1_t, I1_t, I1_t, BF16, BF16, BF16, BF16, ck::Tuple<>, ck::Tuple<>, MaskDisabled_t>,
-    std::tuple<I2_t, I1_t, I1_t, I1_t, I1_t, BF16, BF16, BF16, BF16, ck::Tuple<>, ck::Tuple<>, MaskOutUpperTriangle_t>
+    std::tuple<I2_t, I1_t, I1_t, I1_t, I1_t, BF16, BF16, BF16, BF16, ck::Tuple<>, ck::Tuple<>, MaskUpperTriangleFromTopLeft_t>
     >;
 // clang-format on
 

test/batched_gemm_softmax_gemm_permute/test_batched_gemm_softmax_gemm_permute_fp16.cpp

@@ -15,13 +15,14 @@ using I2_t = ck::Number<2>;
 
 using MaskDisabled_t =
     ck::integral_constant<MaskingSpecialization, MaskingSpecialization::MaskDisabled>;
-using MaskOutUpperTriangle_t =
-    ck::integral_constant<MaskingSpecialization, MaskingSpecialization::MaskOutUpperTriangle>;
+using MaskUpperTriangleFromTopLeft_t =
+    ck::integral_constant<MaskingSpecialization,
+                          MaskingSpecialization::MaskUpperTriangleFromTopLeft>;
 
 // clang-format off
 using KernelTypes = ::testing::Types<
     std::tuple<I2_t, I1_t, I1_t, I1_t, I1_t, F16, F16, F16, F16, ck::Tuple<>, ck::Tuple<>, MaskDisabled_t>,
-    std::tuple<I2_t, I1_t, I1_t, I1_t, I1_t, F16, F16, F16, F16, ck::Tuple<>, ck::Tuple<>, MaskOutUpperTriangle_t>
+    std::tuple<I2_t, I1_t, I1_t, I1_t, I1_t, F16, F16, F16, F16, ck::Tuple<>, ck::Tuple<>, MaskUpperTriangleFromTopLeft_t>
     >;
 // clang-format on
 

test/batched_gemm_softmax_gemm_permute/test_batched_gemm_softmax_gemm_permute_util.hpp

@@ -174,7 +174,7 @@ struct DeviceInstanceWrapper_G2M1N1K1O1_TNTT_FP16_M128_N128_K32_O128
             2,              // CShuffleNXdlPerWavePerShuffle
             S<1, 32, 1, 8>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock
             8,              // CShuffleBlockTransferScalarPerVector_NPerBlock
-            MaskingSpecialization::MaskOutUpperTriangle>; // MaskOutUpperTriangle
+            MaskingSpecialization::MaskUpperTriangleFromTopLeft>; // MaskUpperTriangleFromTopLeft
 
     bool IsSupported(int M, int N, int K, int O)
     {
@@ -315,7 +315,7 @@ struct DeviceInstanceWrapper_G2M1N1K1O1_TNTT_BF16_M128_N128_K32_O128
             2,              // CShuffleNXdlPerWavePerShuffle
             S<1, 32, 1, 8>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock
             8,              // CShuffleBlockTransferScalarPerVector_NPerBlock
-            MaskingSpecialization::MaskOutUpperTriangle>; // MaskOutUpperTriangle
+            MaskingSpecialization::MaskUpperTriangleFromTopLeft>; // MaskUpperTriangleFromTopLeft
 
     bool IsSupported(int M, int N, int K, int O)
     {