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

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

Merge pull request #811 from ROCmSoftwarePlatform/attn-train-develop-qloop-mask
Add another mask(upper tringle from bottom right) to flash attetion
2018bd28 · Dan Yao · GitHub · 120760d6 · a822937a · 2018bd28
Unverified Commit 2018bd28 authored Jul 28, 2023 by Dan Yao Committed by GitHub Jul 28, 2023
7 changed files
--- a/profiler/include/profiler/profile_batched_gemm_softmax_gemm_permute_impl.hpp
+++ b/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();
        });


--- a/test/batched_gemm_softmax_gemm_permute/test_batched_gemm_bias_softmax_gemm_permute_bf16.cpp
+++ b/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


--- a/test/batched_gemm_softmax_gemm_permute/test_batched_gemm_bias_softmax_gemm_permute_fp16.cpp
+++ b/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


--- a/test/batched_gemm_softmax_gemm_permute/test_batched_gemm_bias_softmax_gemm_permute_util.hpp
+++ b/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)
    {

--- a/test/batched_gemm_softmax_gemm_permute/test_batched_gemm_softmax_gemm_permute_bf16.cpp
+++ b/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


--- a/test/batched_gemm_softmax_gemm_permute/test_batched_gemm_softmax_gemm_permute_fp16.cpp
+++ b/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


--- a/test/batched_gemm_softmax_gemm_permute/test_batched_gemm_softmax_gemm_permute_util.hpp
+++ b/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)
    {