/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/fmha/kernel/fmha_fwd_v3_kernel.hpp Source File#
fmha_fwd_v3_kernel.hpp
Go to the documentation of this file.
Definition: cluster_descriptor.hpp:13
constexpr CK_TILE_DEVICE auto make_null_tile_window(const WindowLengths &window_lengths)
Definition: null_tile_window.hpp:66
constexpr CK_TILE_HOST_DEVICE auto integer_divide_ceil(X x, Y y)
Definition: math.hpp:149
constexpr CK_TILE_HOST_DEVICE auto pad_tensor_view(const TensorView &tensor_view, const TileLengths &tile_lengths, DoPads)
Definition: tensor_view.hpp:530
remove_cv_t< std::remove_reference_t< T > > remove_cvref_t
Definition: type_traits.hpp:21
constexpr CK_TILE_DEVICE auto make_tile_window(null_tensor_view, const WindowLengths &window_lengths, const multi_index< WindowLengths::size()> &, Ts &&...)
Definition: null_tile_window.hpp:75
constexpr CK_TILE_HOST_DEVICE auto make_tuple(Xs &&... xs)
Definition: tuple.hpp:360
@ MASK_FROM_TOP_LEFT
typename conditional< predicate, X, Y >::type conditional_t
Definition: functional.hpp:115
Definition: fmha_fwd_v3_kernel.hpp:97
ck_tile::index_t batch_stride_o
Definition: fmha_fwd_v3_kernel.hpp:101
ck_tile::index_t batch_stride_q
Definition: fmha_fwd_v3_kernel.hpp:98
ck_tile::index_t batch_stride_k
Definition: fmha_fwd_v3_kernel.hpp:99
ck_tile::index_t batch_stride_v
Definition: fmha_fwd_v3_kernel.hpp:100
Definition: fmha_fwd_v3_kernel.hpp:51
ck_tile::index_t nhead_stride_v
Definition: fmha_fwd_v3_kernel.hpp:75
ck_tile::index_t nhead_stride_q
Definition: fmha_fwd_v3_kernel.hpp:73
ck_tile::index_t nhead_ratio_qk
Definition: fmha_fwd_v3_kernel.hpp:65
ck_tile::index_t stride_o
Definition: fmha_fwd_v3_kernel.hpp:71
ck_tile::index_t seqlen_k
Definition: fmha_fwd_v3_kernel.hpp:58
ck_tile::index_t nhead_stride_o
Definition: fmha_fwd_v3_kernel.hpp:76
ck_tile::index_t hdim_q
Definition: fmha_fwd_v3_kernel.hpp:59
ck_tile::index_t seqlen_q
Definition: fmha_fwd_v3_kernel.hpp:57
ck_tile::index_t stride_v
Definition: fmha_fwd_v3_kernel.hpp:70
const void * q_ptr
Definition: fmha_fwd_v3_kernel.hpp:52
float scale_s
Definition: fmha_fwd_v3_kernel.hpp:66
const void * v_ptr
Definition: fmha_fwd_v3_kernel.hpp:54
void * o_ptr
Definition: fmha_fwd_v3_kernel.hpp:55
ck_tile::index_t stride_q
Definition: fmha_fwd_v3_kernel.hpp:68
ck_tile::index_t num_head_q
Definition: fmha_fwd_v3_kernel.hpp:62
const void * k_ptr
Definition: fmha_fwd_v3_kernel.hpp:53
ck_tile::index_t hdim_v
Definition: fmha_fwd_v3_kernel.hpp:60
ck_tile::index_t nhead_stride_k
Definition: fmha_fwd_v3_kernel.hpp:74
ck_tile::index_t stride_k
Definition: fmha_fwd_v3_kernel.hpp:69
Definition: fmha_fwd_v3_kernel.hpp:87
void * lse_ptr
Definition: fmha_fwd_v3_kernel.hpp:88
ck_tile::index_t nhead_stride_lse
Definition: fmha_fwd_v3_kernel.hpp:89
ck_tile::index_t batch_stride_lse
Definition: fmha_fwd_v3_kernel.hpp:90
Definition: fmha_fwd_v3_kernel.hpp:44
Definition: fmha_fwd_v3_kernel.hpp:108
const int32_t * seqlen_k_ptr
Definition: fmha_fwd_v3_kernel.hpp:111
const int32_t * seqstart_k_ptr
Definition: fmha_fwd_v3_kernel.hpp:110
const int32_t * seqstart_q_ptr
Definition: fmha_fwd_v3_kernel.hpp:109
Definition: fmha_fwd_v3_kernel.hpp:80
ck_tile::index_t window_size_left
Definition: fmha_fwd_v3_kernel.hpp:82
ck_tile::index_t window_size_right
Definition: fmha_fwd_v3_kernel.hpp:82
ck_tile::GenericAttentionMaskEnum mask_type
Definition: fmha_fwd_v3_kernel.hpp:83
Definition: fmha_fwd_v3_kernel.hpp:17
static constexpr bool kPadSeqLenQ
Definition: fmha_fwd_v3_kernel.hpp:32
ck_tile::remove_cvref_t< typename FmhaPipeline::QDataType > QDataType
Definition: fmha_fwd_v3_kernel.hpp:24
CK_TILE_DEVICE void operator()(Kargs kargs) const
Definition: fmha_fwd_v3_kernel.hpp:307
static constexpr bool kPadSeqLenK
Definition: fmha_fwd_v3_kernel.hpp:33
static constexpr CK_TILE_HOST_DEVICE ck_tile::index_t GetSmemSize()
Definition: fmha_fwd_v3_kernel.hpp:302
static constexpr CK_TILE_HOST auto GridSize(ck_tile::index_t batch_size_, ck_tile::index_t nhead_, ck_tile::index_t seqlen_q_, ck_tile::index_t hdim_v_)
Definition: fmha_fwd_v3_kernel.hpp:258
static constexpr ck_tile::index_t kBlockSize
Definition: fmha_fwd_v3_kernel.hpp:20
ck_tile::remove_cvref_t< typename FmhaPipeline::ODataType > ODataType
Definition: fmha_fwd_v3_kernel.hpp:28
ck_tile::remove_cvref_t< typename FmhaPipeline::KDataType > KDataType
Definition: fmha_fwd_v3_kernel.hpp:25
static constexpr bool kPadHeadDimV
Definition: fmha_fwd_v3_kernel.hpp:35
static constexpr CK_TILE_HOST auto BlockSize()
Definition: fmha_fwd_v3_kernel.hpp:300
ck_tile::remove_cvref_t< FmhaPipeline_ > FmhaPipeline
Definition: fmha_fwd_v3_kernel.hpp:18
static constexpr bool kHasMask
Definition: fmha_fwd_v3_kernel.hpp:39
std::conditional_t< kIsGroupMode, FmhaFwdGroupModeKargs, FmhaFwdBatchModeKargs > Kargs
Definition: fmha_fwd_v3_kernel.hpp:114
static constexpr CK_TILE_HOST std::enable_if_t< Cond, Kargs > MakeKargs(const void *q_ptr, const void *k_ptr, const void *v_ptr, void *lse_ptr, void *o_ptr, const void *seqstart_q_ptr, const void *seqstart_k_ptr, const void *seqlen_k_ptr, ck_tile::index_t hdim_q, ck_tile::index_t hdim_v, ck_tile::index_t num_head_q, ck_tile::index_t nhead_ratio_qk, float scale_s, ck_tile::index_t stride_q, ck_tile::index_t stride_k, ck_tile::index_t stride_v, ck_tile::index_t stride_o, ck_tile::index_t nhead_stride_q, ck_tile::index_t nhead_stride_k, ck_tile::index_t nhead_stride_v, ck_tile::index_t nhead_stride_lse, ck_tile::index_t nhead_stride_o, ck_tile::index_t window_size_left, ck_tile::index_t window_size_right, ck_tile::index_t mask_type)
Definition: fmha_fwd_v3_kernel.hpp:192
ck_tile::remove_cvref_t< EpiloguePipeline_ > EpiloguePipeline
Definition: fmha_fwd_v3_kernel.hpp:19
static constexpr bool kPadHeadDimQ
Definition: fmha_fwd_v3_kernel.hpp:34
ck_tile::remove_cvref_t< typename FmhaPipeline::VDataType > VDataType
Definition: fmha_fwd_v3_kernel.hpp:26
static constexpr ck_tile::index_t kBlockPerCu
Definition: fmha_fwd_v3_kernel.hpp:21
static constexpr CK_TILE_HOST std::enable_if_t< Cond, Kargs > MakeKargs(const void *q_ptr, const void *k_ptr, const void *v_ptr, void *lse_ptr, void *o_ptr, ck_tile::index_t seqlen_q, ck_tile::index_t seqlen_k, ck_tile::index_t hdim_q, ck_tile::index_t hdim_v, ck_tile::index_t num_head_q, ck_tile::index_t nhead_ratio_qk, float scale_s, ck_tile::index_t stride_q, ck_tile::index_t stride_k, ck_tile::index_t stride_v, ck_tile::index_t stride_o, ck_tile::index_t nhead_stride_q, ck_tile::index_t nhead_stride_k, ck_tile::index_t nhead_stride_v, ck_tile::index_t nhead_stride_lse, ck_tile::index_t nhead_stride_o, ck_tile::index_t batch_stride_q, ck_tile::index_t batch_stride_k, ck_tile::index_t batch_stride_v, ck_tile::index_t batch_stride_lse, ck_tile::index_t batch_stride_o, ck_tile::index_t window_size_left, ck_tile::index_t window_size_right, ck_tile::index_t mask_type)
Definition: fmha_fwd_v3_kernel.hpp:118
static constexpr bool kStoreLSE
Definition: fmha_fwd_v3_kernel.hpp:36
ck_tile::remove_cvref_t< typename FmhaPipeline::SaccDataType > SaccDataType
Definition: fmha_fwd_v3_kernel.hpp:29
ck_tile::remove_cvref_t< typename FmhaPipeline::LSEDataType > LSEDataType
Definition: fmha_fwd_v3_kernel.hpp:27
static constexpr bool kIsGroupMode
Definition: fmha_fwd_v3_kernel.hpp:31
ck_tile::remove_cvref_t< typename FmhaPipeline::FmhaMask > FmhaMask
Definition: fmha_fwd_v3_kernel.hpp:38
static constexpr CK_TILE_DEVICE auto GetTileIndex(const Kargs &kargs)
Definition: fmha_fwd_v3_kernel.hpp:270
Definition: integral_constant.hpp:13
Definition: sequence.hpp:49