[CK_TILE] fmha forward split-kv + combine kernels (#1338)

* FA fwd dropout

* FA bwd

* epilogue reuse

* CMakeLists update

* [CK_TILE] support alibi (#1269)

* add alibi support

* fix code

* update code based on comment

* Support more hdim

* fix fp8 bias

* support seqlen_k=0 case

* remove unused printf

* fix format

---------
Co-authored-by: rocking <ChunYu.Lai@amd.com>

* now fwd/bwd can build

* bwd alibi

* add bwd validation stream_config

* update generated filenames

* update bwd kernel launch

* CK_TILE_HOST_DEVICE in philox

* Transpose -> transpose

* format

* format

* format

* Generate the instance for FA required

* format

* fix error in WarpGemm

* Add num_splits option and dummy split-kv api method

* Generate fmha_fwd_splitkv()

* Add SplitKV kernel codegen logics

* Add SplitKV combine kernel codegen logics

* Fix mismatched return type

* Clean-up code

* Replace sentinel value before storing

* Fix wrong l...

include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_qr_ks_vs_async_default_policy.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qx_ks_vs_custom_policy.hpp"
+
+namespace ck_tile {
+
+// This pipeline is qkv all located in LDS
+using BlockFmhaFwdSplitKVPipelineQRKSVSAsyncDefaultPolicy =
+    BlockFmhaPipelineQXKSVSCustomPolicy</* QLoadOnce = */ true,
+                                        /* AsyncCopyK = */ true,
+                                        /* AsyncCopyV = */ false,
+                                        /* NumPrefetchK = */ 3,
+                                        /* NumPrefetchV = */ 3>;
+
+} // namespace ck_tile

include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_qr_ks_vs_default_policy.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qx_ks_vs_custom_policy.hpp"
+
+namespace ck_tile {
+
+// This pipeline is qkv all located in LDS
+using BlockFmhaFwdSplitKVPipelineQRKSVSDefaultPolicy =
+    BlockFmhaPipelineQXKSVSCustomPolicy</* QLoadOnce = */ true,
+                                        /* AsyncCopyK = */ false,
+                                        /* AsyncCopyV = */ false,
+                                        /* NumPrefetchK = */ 1,
+                                        /* NumPrefetchV = */ 1>;
+
+} // namespace ck_tile

include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp

@@ -54,4 +54,69 @@ struct BlockFmhaPipelineProblem
     static constexpr index_t kBlockPerCu    = Traits::kBlockPerCu;
 };
 
+template <typename QDataType,
+          typename KDataType,
+          typename VDataType,
+          typename SaccDataType,
+          typename SMPLComputeDataType,
+          typename BiasDataType,
+          typename RandValOutputDataType,
+          typename LSEDataType,
+          typename PDataType,
+          typename OaccDataType,
+          typename ODataType,
+          typename BlockFmhaShape,
+          bool kIsGroupMode,
+          typename FmhaMask,
+          typename Traits>
+struct BlockFmhaFwdSplitKVPipelineProblem : BlockFmhaPipelineProblem<QDataType,
+                                                                     KDataType,
+                                                                     VDataType,
+                                                                     SaccDataType,
+                                                                     SMPLComputeDataType,
+                                                                     BiasDataType,
+                                                                     RandValOutputDataType,
+                                                                     LSEDataType,
+                                                                     PDataType,
+                                                                     OaccDataType,
+                                                                     ODataType,
+                                                                     BlockFmhaShape,
+                                                                     kIsGroupMode,
+                                                                     FmhaMask,
+                                                                     Traits>
+{
+    static constexpr bool kHasUnevenSplits = kIsGroupMode || Traits::kHasUnevenSplits;
+};
+
+template <typename LSEDataType_,
+          typename OaccDataType_,
+          typename ODataType_,
+          index_t HeadDimV_,
+          index_t kM0_,
+          index_t kN1_,
+          bool kIsGroupMode_,
+          typename Traits_>
+struct BlockFmhaSplitKVCombinePipelineProblem
+{
+    using LSEDataType  = remove_cvref_t<LSEDataType_>;
+    using OaccDataType = remove_cvref_t<OaccDataType_>;
+    using ODataType    = remove_cvref_t<ODataType_>;
+    using Traits       = remove_cvref_t<Traits_>;
+
+    static constexpr index_t kBlockSize = 256;
+    static constexpr bool kIsGroupMode  = kIsGroupMode_;
+
+    static constexpr index_t kHeadDimV = HeadDimV_;
+    static constexpr index_t kM0       = kM0_;
+    static constexpr index_t kN1       = kN1_;
+
+    // attributes from traits
+    static constexpr bool kPadSeqLenQ       = Traits::kPadSeqLenQ;
+    static constexpr bool kPadHeadDimV      = Traits::kPadHeadDimV;
+    static constexpr bool kStoreLSE         = Traits::kStoreLSE;
+    static constexpr bool kDoFp8StaticQuant = Traits::kDoFp8StaticQuant;
+    static constexpr index_t kBlockPerCu    = Traits::kBlockPerCu;
+    static constexpr index_t kMaxSplits     = Traits::kMaxSplits;
+};
+
 } // namespace ck_tile

include/ck_tile/ops/fmha/pipeline/tile_fmha_traits.hpp

@@ -32,6 +32,50 @@ struct TileFmhaTraits
     static constexpr index_t kBlockPerCu    = kBlockPerCu_;
 };
 
+template <bool kPadSeqLenQ /* padding for seqlen_q */,
+          bool kPadSeqLenK /* padding for seqlen_k */,
+          bool kPadHeadDimQ /* paddding for hdim_q */,
+          bool kPadHeadDimV /* paddding for hdim_v */,
+          BlockAttentionBiasEnum BiasEnum,
+          bool kHasBiasGrad,
+          bool kStoreLSE,
+          bool kHasDropout,
+          bool kDoFp8StaticQuant,
+          bool kHasUnevenSplits_ = true,
+          index_t kBlockPerCu    = -1 /* overwrite occupancy if not -1 */>
+struct TileFmhaFwdSplitKVTraits : TileFmhaTraits<kPadSeqLenQ,
+                                                 kPadSeqLenK,
+                                                 kPadHeadDimQ,
+                                                 kPadHeadDimV,
+                                                 BiasEnum,
+                                                 kHasBiasGrad,
+                                                 kStoreLSE,
+                                                 kHasDropout,
+                                                 kDoFp8StaticQuant,
+                                                 kBlockPerCu>
+{
+    // determine if some split (length) is not divisible by tile size
+    static constexpr bool kHasUnevenSplits = kHasUnevenSplits_;
+};
+
+template <bool kPadSeqLenQ_ /* padding for seqlen_q */,
+          bool kPadHeadDimV_ /* paddding for hdim_v */,
+          bool kStoreLSE_,
+          bool kDoFp8StaticQuant_,
+          index_t kLogMaxSplits_,
+          index_t kBlockPerCu_ = -1 /* overwrite occupancy if not -1 */>
+struct TileFmhaFwdSplitKVCombineTraits
+{
+    static constexpr bool kPadSeqLenQ       = kPadSeqLenQ_;
+    static constexpr bool kPadHeadDimV      = kPadHeadDimV_;
+    static constexpr bool kStoreLSE         = kStoreLSE_;
+    static constexpr bool kDoFp8StaticQuant = kDoFp8StaticQuant_;
+
+    static constexpr index_t kMaxSplits = (1 << kLogMaxSplits_);
+    static_assert(kMaxSplits <= get_warp_size() || kMaxSplits % get_warp_size() == 0);
+    static constexpr index_t kBlockPerCu = kBlockPerCu_;
+};
+
 template <bool kPadSeqLenQ_ /* padding for seqlen_q */,
           bool kPadHeadDimV_ /* paddding for hdim_v */,
           index_t kBlockPerCu_ = 2 /* hint to occupancy */>