/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp Source File
Go to the documentation of this file.
 // Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 // SPDX-License-Identifier: MIT
  
 #pragma once
  
 #include "ck_tile/core.hpp"
 #include "ck_tile/ops/common.hpp"
 #include "ck_tile/ops/fmha/block/block_attention_bias_enum.hpp"
 #include "ck_tile/ops/fmha/block/block_attention_kvcache_layout_enum.hpp"
 #include "ck_tile/ops/fmha/block/block_attention_quant_scale_enum.hpp"
 #include "ck_tile/ops/fmha/block/variants.hpp"
  
 #include <string>
 #include <type_traits>
 #include <utility>
 #include <variant>
  
 // S[seqlen_q, seqlen_k] = Q[seqlen_q, hdim_q] @ K[seqlen_k, hdim_q]
 // S'[seqlen_q, seqlen_k] = S[seqlen_q, seqlen_k] * Scale[1]
 // S''[seqlen_q, seqlen_k] = S'[seqlen_q, seqlen_k] + Bias[seqlen_q, seqlen_k]
 // P[seqlen_q, seqlen_k] = Softmax(S''[seqlen_q, seqlen_k])
 // O[seqlen_q, hdim_v] = P[seqlen_q, seqlen_k] @ V^T[hdim_v, seqlen_k]
  
 namespace ck_tile {
  
 template <typename FmhaPipeline_, typename EpiloguePipeline_>
 struct FmhaBatchPrefillWithPagedKVCacheKernel
 {
     using FmhaPipeline                           = ck_tile::remove_cvref_t<FmhaPipeline_>;
     using EpiloguePipeline                       = ck_tile::remove_cvref_t<EpiloguePipeline_>;
     static constexpr ck_tile::index_t kBlockSize = FmhaPipeline::kBlockSize;
  
     static constexpr ck_tile::index_t kBlockPerCu = FmhaPipeline::kBlockPerCu;
     static_assert(kBlockPerCu > 0);
     static constexpr ck_tile::index_t kBlockPerCuInput = FmhaPipeline::Problem::kBlockPerCu;
  
     using QDataType    = ck_tile::remove_cvref_t<typename FmhaPipeline::QDataType>;
     using KDataType    = ck_tile::remove_cvref_t<typename FmhaPipeline::KDataType>;
     using VDataType    = ck_tile::remove_cvref_t<typename FmhaPipeline::VDataType>;
     using PDataType    = ck_tile::remove_cvref_t<typename FmhaPipeline::PDataType>;
     using BiasDataType = ck_tile::remove_cvref_t<typename FmhaPipeline::BiasDataType>;
     using RandValOutputDataType =
         ck_tile::remove_cvref_t<typename FmhaPipeline::RandValOutputDataType>;
     using LSEDataType  = ck_tile::remove_cvref_t<typename FmhaPipeline::LSEDataType>;
     using ODataType    = ck_tile::remove_cvref_t<typename FmhaPipeline::ODataType>;
     using SaccDataType = ck_tile::remove_cvref_t<typename FmhaPipeline::SaccDataType>;
  
     using VLayout = ck_tile::remove_cvref_t<typename FmhaPipeline::VLayout>;
  
     static constexpr bool kIsGroupMode      = FmhaPipeline::kIsGroupMode;
     static constexpr bool kPadSeqLenQ       = FmhaPipeline::kPadSeqLenQ;
     static constexpr bool kPadSeqLenK       = FmhaPipeline::kPadSeqLenK;
     static constexpr bool kPadHeadDimQ      = FmhaPipeline::kPadHeadDimQ;
     static constexpr bool kPadHeadDimV      = FmhaPipeline::kPadHeadDimV;
     static constexpr bool kHasLogitsSoftCap = FmhaPipeline::kHasLogitsSoftCap;
     static constexpr auto BiasEnum          = FmhaPipeline::BiasEnum;
     static constexpr bool kStoreLSE         = FmhaPipeline::kStoreLSE;
     static constexpr bool kHasDropout       = FmhaPipeline::kHasDropout;
     static constexpr auto QScaleEnum        = FmhaPipeline::Problem::QScaleEnum;
     static constexpr auto kKVMemoryLayout   = FmhaPipeline::Problem::kKVMemoryLayout;
     static constexpr auto kKVLookupTable    = FmhaPipeline::Problem::kKVLookupTable;
     static constexpr index_t kPageBlockSize = FmhaPipeline::kPageBlockSize;
     static constexpr index_t kVectorSize    = FmhaPipeline::kVectorSize;
     using AttentionVariant = ck_tile::remove_cvref_t<typename FmhaPipeline::AttentionVariant>;
     using FmhaMask         = ck_tile::remove_cvref_t<typename FmhaPipeline::FmhaMask>;
     static constexpr bool kHasMask = FmhaMask::IsMasking;
  
     static constexpr bool kUseAsyncCopy = FmhaPipeline::Policy::AsyncCopy;
     template <ck_tile::index_t I> // to avoid duplicated base class prblem, introduce an template
                                   // arg
     struct FmhaFwdEmptyKargs
     {
     };
  
     // kargs use aggregate initializer, so no constructor will provided
     // use inheritance to minimize karg size
     // user need to use MakeKargs() function to create kargs.
     struct SglangPageTableKargs
     {
         const int32_t* kv_indptr;
         const int32_t* kv_page_indices;
         const int32_t* kv_last_page_lens;
     };
  
     struct VllmPageTableKargs
     {
         const int32_t* block_table_ptr;
         ck_tile::index_t batch_stride_block_table;
         const int32_t* seqlen_k_ptr;
     };
  
     using PageBlockTableKargs =
         std::conditional_t<kKVLookupTable ==
                                BlockAttentionKVCacheLookupTableEnum::SGLANG_PAGE_TABLE_1D,
                            SglangPageTableKargs,
                            VllmPageTableKargs>;
  
     struct FmhaFwdCommonKargs
     {
         const void* q_ptr;
         const void* k_ptr;
         const void* v_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;
         // for MQA/GQA, nhead could be different. This parameter is nhead_q / nhead_k
         // if this param is larger than 1, indicate MQA/GQA case
         ck_tile::index_t nhead_ratio_qk;
  
         int32_t num_total_pages;
         ck_tile::index_t page_block_size;
         PageBlockTableKargs page_table;
  
         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_o;
     };
  
     struct FmhaFwdLogitsSoftCapKargs
     {
         FmhaFwdLogitsSoftCapKargs() = default;
  
         void init_logits_soft_cap(float logits_soft_cap_)
         {
             if(0 < logits_soft_cap_)
             {
                 logits_soft_cap     = logits_soft_cap_;
                 logits_soft_cap_rcp = 1.f / logits_soft_cap;
             }
             else
             {
                 logits_soft_cap     = 0.f;
                 logits_soft_cap_rcp = 0.f;
             }
         }
  
         float logits_soft_cap;
         float logits_soft_cap_rcp;
     };
  
     struct FmhaFwdCommonBiasKargs
     {
         const void* bias_ptr               = nullptr;
         ck_tile::index_t stride_bias       = 0;
         ck_tile::index_t nhead_stride_bias = 0;
     };
  
     struct FmhaFwdBatchModeBiasKargs : FmhaFwdCommonBiasKargs
     {
         ck_tile::index_t batch_stride_bias = 0;
     };
  
     struct FmhaFwdAlibiKargs
     {
         // alibi is batch*nhead*1, no matter in batch/group mode, they are the same
         const void* alibi_slope_ptr;
         ck_tile::index_t alibi_slope_stride; // stride in batch, or 0 for all batch share same slope
     };
  
     struct FmhaFwdMaskKargs
     {
         // ck_tile::index_t window_size_left, window_size_right;
         ck_tile::index_t window_size_left, window_size_right, sink_size;
         ck_tile::GenericAttentionMaskEnum mask_type;
     };
  
     struct FmhaFwdCommonLSEKargs
     {
         void* lse_ptr                     = nullptr;
         ck_tile::index_t nhead_stride_lse = 0;
         ck_tile::index_t batch_stride_lse = 0;
     };
  
     struct FmhaFwdCommonQScaleKargs
     {
         const void* q_descale_ptr = nullptr;
         const void* k_descale_ptr = nullptr;
         const void* v_descale_ptr = nullptr;
     };
  
     struct FmhaFwdDropoutSeedOffset
     {
         template <typename T>
         union ValueOrPointer
         {
             T val;
             const T* ptr;
         };
  
         ValueOrPointer<uint64_t> drop_seed;
         ValueOrPointer<uint64_t> drop_offset;
         bool is_drop_seed_offset_from_host;
     };
  
     struct FmhaFwdCommonDropoutKargs : FmhaFwdDropoutSeedOffset
     {
         void init_dropout(float p_drop, uint64_t seed, uint64_t offset)
         {
             float p_undrop = 1.0 - p_drop;
             p_undrop_in_uint8_t =
                 uint8_t(std::floor(p_undrop * std::numeric_limits<uint8_t>::max()));
             rp_undrop = 1.0 / p_undrop;
  
             this->drop_seed.val                 = seed;
             this->drop_offset.val               = offset;
             this->is_drop_seed_offset_from_host = true;
         }
  
         void init_dropout(float p_drop, const uint64_t* seed_ptr, const uint64_t* offset_ptr)
         {
             float p_undrop = 1.0 - p_drop;
             p_undrop_in_uint8_t =
                 uint8_t(std::floor(p_undrop * std::numeric_limits<uint8_t>::max()));
             rp_undrop = 1.0 / p_undrop;
  
             this->drop_seed.ptr                 = seed_ptr;
             this->drop_offset.ptr               = offset_ptr;
             this->is_drop_seed_offset_from_host = false;
         }
  
         float rp_undrop             = 1;
         uint8_t p_undrop_in_uint8_t = std::numeric_limits<uint8_t>::max();
         bool is_store_randval       = false;
         void* rand_val_ptr          = nullptr;
  
         ck_tile::index_t stride_randval       = 0;
         ck_tile::index_t nhead_stride_randval = 0;
     };
  
     struct FmhaFwdBatchModeDropoutKargs : FmhaFwdCommonDropoutKargs
     {
         ck_tile::index_t batch_stride_randval = 0;
     };
  
     struct FmhaFwdBatchModeKargs
         : FmhaFwdCommonKargs,
           std::conditional_t<BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS,
                              FmhaFwdBatchModeBiasKargs,
                              std::conditional_t<BiasEnum == BlockAttentionBiasEnum::ALIBI,
                                                 FmhaFwdAlibiKargs,
                                                 FmhaFwdEmptyKargs<0>>>,
           std::conditional_t<kHasMask, FmhaFwdMaskKargs, FmhaFwdEmptyKargs<1>>,
           std::conditional_t<kStoreLSE, FmhaFwdCommonLSEKargs, FmhaFwdEmptyKargs<2>>,
           std::conditional_t<QScaleEnum == BlockAttentionQuantScaleEnum::PERTENSOR,
                              FmhaFwdCommonQScaleKargs,
                              FmhaFwdEmptyKargs<3>>,
           std::conditional_t<kHasDropout, FmhaFwdBatchModeDropoutKargs, FmhaFwdEmptyKargs<4>>,
           std::conditional_t<kHasLogitsSoftCap, FmhaFwdLogitsSoftCapKargs, FmhaFwdEmptyKargs<5>>
     {
         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_o;
     };
  
     struct FmhaFwdGroupModeKargs
         : FmhaFwdCommonKargs,
           std::conditional_t<BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS,
                              FmhaFwdCommonBiasKargs,
                              std::conditional_t<BiasEnum == BlockAttentionBiasEnum::ALIBI,
                                                 FmhaFwdAlibiKargs,
                                                 FmhaFwdEmptyKargs<0>>>,
           std::conditional_t<kHasMask, FmhaFwdMaskKargs, FmhaFwdEmptyKargs<1>>,
           std::conditional_t<kStoreLSE, FmhaFwdCommonLSEKargs, FmhaFwdEmptyKargs<2>>,
           std::conditional_t<QScaleEnum == BlockAttentionQuantScaleEnum::PERTENSOR,
                              FmhaFwdCommonQScaleKargs,
                              FmhaFwdEmptyKargs<3>>,
           std::conditional_t<kHasDropout, FmhaFwdCommonDropoutKargs, FmhaFwdEmptyKargs<4>>,
           std::conditional_t<kHasLogitsSoftCap, FmhaFwdLogitsSoftCapKargs, FmhaFwdEmptyKargs<5>>
     {
         const int32_t* seqstart_q_ptr;
         ck_tile::index_t batch_stride_k;
         ck_tile::index_t batch_stride_v;
     };
  
     using Kargs = std::conditional_t<kIsGroupMode, FmhaFwdGroupModeKargs, FmhaFwdBatchModeKargs>;
  
     struct BlockIndices
     {
         ck_tile::index_t batch_idx;
         ck_tile::index_t qo_head_idx;
         ck_tile::index_t kv_head_idx;
     };
  
     template <bool Cond = !kIsGroupMode>
     CK_TILE_HOST static constexpr std::enable_if_t<Cond, Kargs>
     MakeKargs(const void* q_ptr,
               const void* k_ptr,
               const void* v_ptr,
               const void* bias_ptr,
               const void* q_descale_ptr,
               const void* k_descale_ptr,
               const void* v_descale_ptr,
               void* rand_val_ptr,
               void* lse_ptr,
               void* o_ptr,
               ck_tile::index_t seqlen_q,
               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,
               int32_t num_total_pages,
               ck_tile::index_t page_block_size,
               const PageBlockTableKargs& page_table,
               float scale_s,
               [[maybe_unused]] float scale_p,
               [[maybe_unused]] float scale_o,
               float logits_soft_cap,
               ck_tile::index_t stride_q,
               ck_tile::index_t stride_k,
               ck_tile::index_t stride_v,
               ck_tile::index_t stride_bias,
               ck_tile::index_t stride_randval,
               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_bias,
               ck_tile::index_t nhead_stride_randval,
               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_bias,
               ck_tile::index_t batch_stride_randval,
               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 sink_size,
               ck_tile::index_t mask_type,
               float p_drop,
               bool s_randval,
               std::variant<std::pair<uint64_t, uint64_t>, std::pair<const void*, const void*>>
                   drop_seed_offset)
     {
         Kargs kargs{{q_ptr,
                      k_ptr,
                      v_ptr,
                      o_ptr,
                      seqlen_q,
                      -1,
                      hdim_q,
                      hdim_v,
                      num_head_q,
                      nhead_ratio_qk,
                      num_total_pages,
                      page_block_size,
                      page_table,
 #if CK_TILE_FMHA_FWD_FAST_EXP2
                      static_cast<float>(scale_s * ck_tile::log2e_v<>),
 #else
                      scale_s,
 #endif
                      stride_q,
                      stride_k,
                      stride_v,
                      stride_o,
                      nhead_stride_q,
                      nhead_stride_k,
                      nhead_stride_v,
                      nhead_stride_o}, // args for common karg
                     {},               // placeholder for bias
                     {},               // placeholder for mask
                     {},               // placeholder for lse
                     {},               // placeholder for qscale
                     {},               // placeholder for dropout
                     {},               // placeholder for logits_soft_cap
                     batch_stride_q,
                     batch_stride_k,
                     batch_stride_v,
                     batch_stride_o};
  
         if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
         {
             kargs.bias_ptr          = bias_ptr;
             kargs.stride_bias       = stride_bias;
             kargs.nhead_stride_bias = nhead_stride_bias;
             kargs.batch_stride_bias = batch_stride_bias;
         }
         else if constexpr(BiasEnum == BlockAttentionBiasEnum::ALIBI)
         {
             kargs.alibi_slope_ptr    = bias_ptr;
             kargs.alibi_slope_stride = stride_bias;
         }
         if constexpr(kHasMask)
         {
             kargs.window_size_left  = window_size_left;
             kargs.window_size_right = window_size_right;
             kargs.sink_size         = sink_size;
             kargs.mask_type         = static_cast<ck_tile::GenericAttentionMaskEnum>(mask_type);
         }
         if constexpr(kStoreLSE)
         {
             kargs.lse_ptr          = lse_ptr;
             kargs.nhead_stride_lse = nhead_stride_lse;
             kargs.batch_stride_lse = batch_stride_lse;
         }
         if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::PERTENSOR)
         {
             kargs.q_descale_ptr = q_descale_ptr;
             kargs.k_descale_ptr = k_descale_ptr;
             kargs.v_descale_ptr = v_descale_ptr;
         }
         if constexpr(kHasDropout)
         {
             if(drop_seed_offset.index() == 0) // seed & offset come from host
             {
                 const auto& [seed, offset] = std::get<0>(drop_seed_offset);
                 kargs.init_dropout(p_drop, seed, offset);
             }
             else // seed & offset come from device
             {
                 const auto& [seed_ptr, offset_ptr] = std::get<1>(drop_seed_offset);
                 kargs.init_dropout(p_drop,
                                    reinterpret_cast<const uint64_t*>(seed_ptr),
                                    reinterpret_cast<const uint64_t*>(offset_ptr));
             }
  
             kargs.rand_val_ptr         = rand_val_ptr;
             kargs.stride_randval       = stride_randval;
             kargs.nhead_stride_randval = nhead_stride_randval;
             kargs.batch_stride_randval = batch_stride_randval;
             kargs.is_store_randval     = s_randval;
         }
         if constexpr(kHasLogitsSoftCap)
         {
             kargs.init_logits_soft_cap(logits_soft_cap);
         }
  
         return kargs;
     }
  
     template <bool Cond = kIsGroupMode>
     CK_TILE_HOST static constexpr std::enable_if_t<Cond, Kargs>
     MakeKargs(const void* q_ptr,
               const void* k_ptr,
               const void* v_ptr,
               const void* bias_ptr,
               const void* q_descale_ptr,
               const void* k_descale_ptr,
               const void* v_descale_ptr,
               void* rand_val_ptr,
               void* lse_ptr,
               void* o_ptr,
               const void* seqstart_q_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,
               int32_t num_total_pages,
               ck_tile::index_t page_block_size,
               const PageBlockTableKargs& page_table,
               float scale_s,
               [[maybe_unused]] float scale_p,
               [[maybe_unused]] float scale_o,
               float logits_soft_cap,
               ck_tile::index_t stride_q,
               ck_tile::index_t stride_k,
               ck_tile::index_t stride_v,
               ck_tile::index_t stride_bias,
               ck_tile::index_t stride_randval,
               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_bias,
               ck_tile::index_t nhead_stride_randval,
               ck_tile::index_t nhead_stride_lse,
               ck_tile::index_t nhead_stride_o,
               ck_tile::index_t batch_stride_k,
               ck_tile::index_t batch_stride_v,
               ck_tile::index_t window_size_left,
               ck_tile::index_t window_size_right,
               ck_tile::index_t sink_size,
               ck_tile::index_t mask_type,
               float p_drop,
               bool s_randval,
               std::variant<std::pair<uint64_t, uint64_t>, std::pair<const void*, const void*>>
                   drop_seed_offset)
     {
         Kargs kargs{{q_ptr,
                      k_ptr,
                      v_ptr,
                      o_ptr,
                      -1, // seqlen will be updated by another pointer
                      -1, //
                      hdim_q,
                      hdim_v,
                      num_head_q,
                      nhead_ratio_qk,
                      num_total_pages,
                      page_block_size,
                      page_table,
 #if CK_TILE_FMHA_FWD_FAST_EXP2
                      static_cast<float>(scale_s * ck_tile::log2e_v<>),
 #else
                      scale_s,
 #endif
                      stride_q,
                      stride_k,
                      stride_v,
                      stride_o,
                      nhead_stride_q,
                      nhead_stride_k,
                      nhead_stride_v,
                      nhead_stride_o}, // args for common karg
                     {},               // placeholder for bias
                     {},               // placeholder for mask
                     {},               // placeholder for lse
                     {},               // placeholder for qscale
                     {},               // placeholder for dropout
                     {},               // placeholder for logits_soft_cap
                     reinterpret_cast<const int32_t*>(seqstart_q_ptr),
                     batch_stride_k,
                     batch_stride_v};
  
         if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
         {
             kargs.bias_ptr          = bias_ptr;
             kargs.stride_bias       = stride_bias;
             kargs.nhead_stride_bias = nhead_stride_bias;
         }
         else if constexpr(BiasEnum == BlockAttentionBiasEnum::ALIBI)
         {
             kargs.alibi_slope_ptr    = bias_ptr;
             kargs.alibi_slope_stride = stride_bias;
         }
         if constexpr(kHasMask)
         {
             kargs.window_size_left  = window_size_left;
             kargs.window_size_right = window_size_right;
             kargs.sink_size         = sink_size;
             kargs.mask_type         = static_cast<ck_tile::GenericAttentionMaskEnum>(mask_type);
         }
         if constexpr(kStoreLSE)
         {
             kargs.lse_ptr          = lse_ptr;
             kargs.nhead_stride_lse = nhead_stride_lse;
         }
         if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::PERTENSOR)
         {
             kargs.q_descale_ptr = q_descale_ptr;
             kargs.k_descale_ptr = k_descale_ptr;
             kargs.v_descale_ptr = v_descale_ptr;
         }
         if constexpr(kHasDropout)
         {
             if(drop_seed_offset.index() == 0) // seed & offset come from host
             {
                 const auto& [seed, offset] = std::get<0>(drop_seed_offset);
                 kargs.init_dropout(p_drop, seed, offset);
             }
             else // seed & offset come from device
             {
                 const auto& [seed_ptr, offset_ptr] = std::get<1>(drop_seed_offset);
                 kargs.init_dropout(p_drop,
                                    reinterpret_cast<const uint64_t*>(seed_ptr),
                                    reinterpret_cast<const uint64_t*>(offset_ptr));
             }
  
             kargs.rand_val_ptr         = rand_val_ptr;
             kargs.stride_randval       = stride_randval;
             kargs.nhead_stride_randval = nhead_stride_randval;
             kargs.is_store_randval     = s_randval;
         }
         if constexpr(kHasLogitsSoftCap)
         {
             kargs.init_logits_soft_cap(logits_soft_cap);
         }
  
         return kargs;
     }
  
     CK_TILE_HOST static constexpr 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_)
     {
         if constexpr(kIsGroupMode)
         {
             // TODO: this may need tuning
             return dim3(nhead_,
                         batch_size_,
                         ck_tile::integer_divide_ceil(seqlen_q_, FmhaPipeline::kM0) *
                             ck_tile::integer_divide_ceil(hdim_v_, FmhaPipeline::kN1));
         }
         else
         {
             // TODO: this may need tuning
             return dim3(ck_tile::integer_divide_ceil(seqlen_q_, FmhaPipeline::kM0) *
                             ck_tile::integer_divide_ceil(hdim_v_, FmhaPipeline::kN1),
                         nhead_,
                         batch_size_);
         }
     }
  
     CK_TILE_DEVICE static constexpr auto GetTileIndex(const Kargs& kargs)
     {
         if constexpr(kIsGroupMode)
         {
             // const index_t num_tile_m0 = seqlen_q / kM0;
             const index_t num_tile_n1 =
                 ck_tile::integer_divide_ceil(kargs.hdim_v, FmhaPipeline::kN1);
  
             const index_t i_block = blockIdx.z;
             const index_t i_nhead = blockIdx.x;
             const index_t i_batch = blockIdx.y;
  
             const auto f = [](index_t dividend, index_t divisor) {
                 index_t quotient = dividend / divisor;
                 index_t modulus  = dividend - quotient * divisor;
                 return ck_tile::make_tuple(quotient, modulus);
             };
  
             const auto [i_tile_m, i_tile_n] = f(i_block, num_tile_n1);
             if constexpr(kHasMask)
             {
                 // assume that num_tile_n1 is always 1
                 return ck_tile::make_tuple(gridDim.z - 1 - i_tile_m, i_tile_n, i_nhead, i_batch);
             }
             else
             {
                 return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch);
             }
         }
         else
         {
             // const index_t num_tile_m0 = seqlen_q / kM0;
             const index_t num_tile_n1 =
                 ck_tile::integer_divide_ceil(kargs.hdim_v, FmhaPipeline::kN1);
  
             const index_t i_block = blockIdx.x;
             const index_t i_nhead = blockIdx.y;
             const index_t i_batch = blockIdx.z;
  
             const auto f = [](index_t dividend, index_t divisor) {
                 index_t quotient = dividend / divisor;
                 index_t modulus  = dividend - quotient * divisor;
                 return ck_tile::make_tuple(quotient, modulus);
             };
  
             const auto [i_tile_m, i_tile_n] = f(i_block, num_tile_n1);
  
             if constexpr(kHasMask)
             {
                 // assume that num_tile_n1 is always 1
                 return ck_tile::make_tuple(gridDim.x - 1 - i_tile_m, i_tile_n, i_nhead, i_batch);
             }
             else
             {
                 return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch);
             }
         }
     }
  
     CK_TILE_HOST static dim3 BlockSize()
     {
         if(is_wave32())
         {
             return dim3(kBlockSize / 2);
         }
         else
         {
             return dim3(kBlockSize);
         }
     }
  
     CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize()
     {
         return ck_tile::max(FmhaPipeline::GetSmemSize(), EpiloguePipeline::GetSmemSize());
     }
  
     CK_TILE_DEVICE void operator()(Kargs kargs) const
     {
         // allocate LDS
         __shared__ char smem_ptr[GetSmemSize()];
  
         // divide problem
         const auto [i_tile_m, i_tile_n, i_nhead, i_batch] = GetTileIndex(kargs);
  
         const index_t i_m0 = amd_wave_read_first_lane(i_tile_m * FmhaPipeline::kM0);
         const index_t i_n1 = amd_wave_read_first_lane(i_tile_n * FmhaPipeline::kN1);
  
         long_index_t batch_offset_q       = 0;
         long_index_t batch_offset_bias    = 0;
         long_index_t batch_offset_randval = 0;
         long_index_t batch_offset_lse     = 0;
         long_index_t batch_offset_o       = 0;
  
         const index_t seqlen_k = [&]() {
             if constexpr(kKVLookupTable ==
                          BlockAttentionKVCacheLookupTableEnum::SGLANG_PAGE_TABLE_1D)
             {
                 const int32_t page_start      = kargs.page_table.kv_indptr[i_batch];
                 const int32_t page_end        = kargs.page_table.kv_indptr[i_batch + 1];
                 const int32_t num_page_blocks = page_end - page_start;
                 const int32_t last_page_len   = [&]() {
                     if constexpr(kPageBlockSize == 1)
                         return static_cast<int32_t>(kPageBlockSize);
                     else
                         return kargs.page_table.kv_last_page_lens[i_batch];
                 }();
                 return num_page_blocks > 0
                            ? static_cast<index_t>((num_page_blocks - 1) * kargs.page_block_size +
                                                   last_page_len)
                            : 0;
             }
             else // BlockAttentionKVCacheLookupTableEnum::VLLM_BLOCK_TABLE_2D
             {
                 if(kargs.page_table.seqlen_k_ptr != nullptr)
                     return static_cast<index_t>(kargs.page_table.seqlen_k_ptr[i_batch]);
                 else
                     return kargs.seqlen_k;
             }
         }();
         const int32_t* page_idx = [&]() {
             if constexpr(kKVLookupTable ==
                          BlockAttentionKVCacheLookupTableEnum::SGLANG_PAGE_TABLE_1D)
             {
                 return kargs.page_table.kv_page_indices + kargs.page_table.kv_indptr[i_batch];
             }
             else // BlockAttentionKVCacheLookupTableEnum::VLLM_BLOCK_TABLE_2D
             {
                 return kargs.page_table.block_table_ptr +
                        static_cast<long_index_t>(i_batch) *
                            kargs.page_table.batch_stride_block_table;
             }
         }();
  
         if constexpr(kIsGroupMode)
         {
             // get starting offset for each batch
             const long_index_t query_start = kargs.seqstart_q_ptr[i_batch];
  
             batch_offset_q = query_start * kargs.stride_q;
  
             if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
             {
                 batch_offset_bias = query_start * kargs.stride_bias;
             }
             if constexpr(kStoreLSE)
             {
                 batch_offset_lse = query_start;
             }
             if constexpr(kHasDropout)
             {
                 batch_offset_randval = query_start * kargs.stride_randval;
             }
             batch_offset_o = query_start * kargs.stride_o;
  
             // get real # queries & # keys under group mode
             kargs.seqlen_q = kargs.seqstart_q_ptr[i_batch + 1] - query_start;
  
             // # of required blocks is different in each groups, terminate unnecessary blocks
             // earlier
             if(kargs.seqlen_q <= i_m0)
             {
                 return;
             }
  
             kargs.seqlen_k = seqlen_k;
         }
         else
         {
             batch_offset_q = static_cast<long_index_t>(i_batch) * kargs.batch_stride_q;
  
             if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
             {
                 batch_offset_bias = static_cast<long_index_t>(i_batch) * kargs.batch_stride_bias;
             }
             if constexpr(kStoreLSE)
             {
                 batch_offset_lse = static_cast<long_index_t>(i_batch) * kargs.batch_stride_lse;
             }
             if constexpr(kHasDropout)
             {
                 batch_offset_randval =
                     static_cast<long_index_t>(i_batch) * kargs.batch_stride_randval;
             }
             batch_offset_o = static_cast<long_index_t>(i_batch) * kargs.batch_stride_o;
  
             kargs.seqlen_k = seqlen_k;
         }
  
         // for simplicity, batch stride we just modify the pointer
         const QDataType* q_ptr = reinterpret_cast<const QDataType*>(kargs.q_ptr) +
                                  static_cast<long_index_t>(i_nhead) * kargs.nhead_stride_q +
                                  batch_offset_q;
         const KDataType* k_ptr =
             reinterpret_cast<const KDataType*>(kargs.k_ptr) +
             static_cast<long_index_t>(i_nhead / kargs.nhead_ratio_qk) * kargs.nhead_stride_k;
         const VDataType* v_ptr =
             reinterpret_cast<const VDataType*>(kargs.v_ptr) +
             static_cast<long_index_t>(i_nhead / kargs.nhead_ratio_qk) * kargs.nhead_stride_v;
         ODataType* o_ptr = reinterpret_cast<ODataType*>(kargs.o_ptr) +
                            static_cast<long_index_t>(i_nhead) * kargs.nhead_stride_o +
                            batch_offset_o;
  
         // Q/K/V DRAM and DRAM window
         const auto q_dram = [&]() {
             const auto q_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                 q_ptr,
                 make_tuple(kargs.seqlen_q, kargs.hdim_q),
                 make_tuple(kargs.stride_q, 1),
                 number<FmhaPipeline::kAlignmentQ>{},
                 number<1>{});
             if constexpr(FmhaPipeline::kQLoadOnce)
             {
                 return pad_tensor_view(
                     q_dram_naive,
                     make_tuple(number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kSubQKHeaddim>{}),
                     sequence<kPadSeqLenQ, kPadHeadDimQ>{});
             }
             else
             {
                 return pad_tensor_view(
                     q_dram_naive,
                     make_tuple(number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kK0>{}),
                     sequence<kPadSeqLenQ, kPadHeadDimQ>{});
             }
         }();
         const auto k_dram = [&]() {
             if constexpr(kKVMemoryLayout ==
                          BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT)
             {
                 // Vectorized K Layout: [NumPages, D/kVectorSize, S, kVectorSize]
                 // Logical View for Pipeline: (TotalSeqK, D)
  
                 // Define the naive physical view with 4D shape: (NumPages, HeadDim/kVectorSize,
                 // PageBlockSize, kVectorSize)
                 //    Strides: (BatchStride, PageBlockSize*kVectorSize, kVectorSize, 1)
                 const auto k_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                     k_ptr,
                     make_tuple(kargs.num_total_pages,
                                kargs.hdim_q / kVectorSize,
                                kargs.page_block_size,
                                kVectorSize),
                     make_tuple(
                         kargs.batch_stride_k, kargs.page_block_size * kVectorSize, kVectorSize, 1),
                     number<FmhaPipeline::kAlignmentK>{},
                     number<1>{});
  
                 // Merge to (TotalSeqK, D) in a single transform:
                 // physical (Page, D/vec, S, vec) -> logical (TotalSeqK, D)
                 auto k_dram_2d = transform_tensor_view(
                     k_dram_naive,
                     make_tuple(make_merge_transform(make_tuple(kargs.num_total_pages,
                                                                kargs.page_block_size)), // TotalSeqK
                                make_merge_transform(
                                    make_tuple(static_cast<int32_t>(kargs.hdim_q / kVectorSize),
                                               static_cast<int32_t>(kVectorSize)))), // D
                     make_tuple(sequence<0, 2>{}, sequence<1, 3>{}),
                     make_tuple(sequence<0>{}, sequence<1>{}));
  
                 constexpr bool kPadSeqLenK_ = kUseAsyncCopy ? kPadSeqLenK : true;
                 return pad_tensor_view(
                     k_dram_2d,
                     make_tuple(number<FmhaPipeline::kN0>{}, number<FmhaPipeline::kK0>{}),
                     sequence<kPadSeqLenK_, kPadHeadDimQ>{});
             }
             else
             {
                 // Linear K Layout: [NumPages, PageSize, NumHeads, HeadDim]
                 // Logical View for Pipeline: (TotalSeqK, D)
                 const auto k_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                     k_ptr,
                     make_tuple(kargs.num_total_pages, kargs.page_block_size, kargs.hdim_q),
                     make_tuple(kargs.batch_stride_k, kargs.stride_k, 1),
                     number<FmhaPipeline::kAlignmentK>{},
                     number<1>{});
  
                 // Merge to (TotalSeqK, D) in a single transform:
                 // physical (Page, S, D) -> logical (TotalSeqK, D)
                 auto k_dram_2d = transform_tensor_view(
                     k_dram_naive,
                     make_tuple(make_merge_transform(
                                    make_tuple(kargs.num_total_pages, kargs.page_block_size)),
                                make_pass_through_transform(kargs.hdim_q)),
                     make_tuple(sequence<0, 1>{}, sequence<2>{}),
                     make_tuple(sequence<0>{}, sequence<1>{}));
  
                 constexpr bool kPadSeqLenK_ = kUseAsyncCopy ? kPadSeqLenK : true;
                 return pad_tensor_view(
                     k_dram_2d,
                     make_tuple(number<FmhaPipeline::kN0>{}, number<FmhaPipeline::kK0>{}),
                     sequence<kPadSeqLenK_, kPadHeadDimQ>{});
             }
         }();
         const auto v_dram = [&]() {
             if constexpr(kKVMemoryLayout ==
                          BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT)
             {
                 // Vectorized V Layout: [NumPages, S/kVectorSize, D, kVectorSize]
                 // Logical View for Pipeline: (D, TotalSeqK) - Transposed for GEMM
  
                 // Define the naive physical view with 4D shape: (NumPages,
                 // PageBlockSize/kVectorSize, HeadDim, kVectorSize)
                 //    Strides: (BatchStride, HeadDim*kVectorSize, kVectorSize, 1)
                 const auto v_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                     v_ptr,
                     make_tuple(kargs.num_total_pages,
                                kargs.page_block_size / kVectorSize,
                                kargs.hdim_v,
                                kVectorSize),
                     make_tuple(kargs.batch_stride_v, kargs.hdim_v * kVectorSize, kVectorSize, 1),
                     number<FmhaPipeline::kAlignmentV>{},
                     number<1>{});
  
                 // Merge to (D, TotalSeqK) in a single transform:
                 // physical (Page, S/vec, D, vec) -> logical (D, TotalSeqK)
                 auto v_dram_final = transform_tensor_view(
                     v_dram_naive,
                     make_tuple(make_pass_through_transform(kargs.hdim_v), // D
                                make_merge_transform(make_tuple(kargs.num_total_pages,
                                                                kargs.page_block_size / kVectorSize,
                                                                kVectorSize))), // TotalSeqK
                     make_tuple(sequence<2>{}, sequence<0, 1, 3>{}),
                     make_tuple(sequence<0>{}, sequence<1>{}));
  
                 constexpr bool kPadSeqLenK_ = kUseAsyncCopy ? kPadSeqLenK : true;
                 return pad_tensor_view(
                     v_dram_final,
                     make_tuple(number<FmhaPipeline::kN1>{}, number<FmhaPipeline::kK1>{}),
                     sequence<kPadHeadDimV, kPadSeqLenK_>{});
             }
             else
             {
                 // Linear V Layout: [NumPages, PageSize, NumHeads, HeadDim]
                 // Logical View for Pipeline: (D, TotalSeqK)
                 const auto v_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                     v_ptr,
                     make_tuple(kargs.num_total_pages, kargs.page_block_size, kargs.hdim_v),
                     make_tuple(kargs.batch_stride_v, kargs.stride_v, 1),
                     number<FmhaPipeline::kAlignmentV>{},
                     number<1>{});
  
                 // Merge to (D, TotalSeqK) in a single transform:
                 // physical (Page, S, D) -> logical (D, TotalSeqK)
                 auto v_dram_final = transform_tensor_view(
                     v_dram_naive,
                     make_tuple(make_pass_through_transform(kargs.hdim_v),
                                make_merge_transform(
                                    make_tuple(kargs.num_total_pages, kargs.page_block_size))),
                     make_tuple(sequence<2>{}, sequence<0, 1>{}),
                     make_tuple(sequence<0>{}, sequence<1>{}));
  
                 constexpr bool kPadSeqLenK_ = kUseAsyncCopy ? kPadSeqLenK : true;
                 return pad_tensor_view(
                     v_dram_final,
                     make_tuple(number<FmhaPipeline::kN1>{}, number<FmhaPipeline::kK1>{}),
                     sequence<kPadHeadDimV, kPadSeqLenK_>{});
             }
         }();
         auto q_dram_window = make_tile_window(
             q_dram,
             [&]() {
                 if constexpr(FmhaPipeline::kQLoadOnce)
                     return make_tuple(number<FmhaPipeline::kM0>{},
                                       number<FmhaPipeline::kSubQKHeaddim>{});
                 else
                     return make_tuple(number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kK0>{});
             }(),
             {i_m0, 0});
  
         auto k_dram_window = make_tile_window(
             k_dram, make_tuple(number<FmhaPipeline::kN0>{}, number<FmhaPipeline::kK0>{}), {0, 0});
  
         auto v_dram_window =
             make_tile_window(v_dram,
                              make_tuple(number<FmhaPipeline::kN1>{}, number<FmhaPipeline::kK1>{}),
                              {i_n1, 0});
         const auto bias_dram_window = [&, i_nhead_ = i_nhead]() {
             constexpr auto bias_dram_window_lengths =
                 make_tuple(number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kN0>{});
             if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
             {
                 const BiasDataType* bias_ptr =
                     reinterpret_cast<const BiasDataType*>(kargs.bias_ptr) +
                     static_cast<long_index_t>(i_nhead_) * kargs.nhead_stride_bias +
                     batch_offset_bias;
  
                 const auto bias_dram = [&]() {
                     const auto bias_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                         bias_ptr,
                         make_tuple(kargs.seqlen_q, kargs.seqlen_k),
                         make_tuple(kargs.stride_bias, 1),
                         number<FmhaPipeline::kAlignmentBias>{},
                         number<1>{});
  
                     return pad_tensor_view(bias_dram_naive,
                                            bias_dram_window_lengths,
                                            sequence<kPadSeqLenQ, kPadSeqLenK>{});
                 }();
  
                 return make_tile_window(bias_dram, bias_dram_window_lengths, {i_m0, 0});
             }
             else
             {
                 return make_null_tile_window(bias_dram_window_lengths);
             }
         }();
  
         // lse
         auto lse_dram_window = [&, i_nhead_ = i_nhead]() {
             constexpr auto lse_dram_window_lengths = make_tuple(number<FmhaPipeline::kM0>{});
             if constexpr(kStoreLSE)
             {
                 LSEDataType* lse_ptr =
                     reinterpret_cast<LSEDataType*>(kargs.lse_ptr) +
                     static_cast<long_index_t>(i_nhead_) * kargs.nhead_stride_lse + batch_offset_lse;
  
                 const auto lse_dram = [&]() {
                     const auto lse_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                         lse_ptr,
                         make_tuple(kargs.seqlen_q),
                         make_tuple(1),
                         number<1>{},
                         number<1>{});
  
                     return pad_tensor_view(
                         lse_dram_naive, lse_dram_window_lengths, sequence<kPadSeqLenQ>{});
                 }();
  
                 return make_tile_window(lse_dram, lse_dram_window_lengths, {i_m0});
             }
             else
             {
                 return make_null_tile_window(lse_dram_window_lengths);
             }
         }();
  
         auto dropout = [&, i_nhead_ = i_nhead, i_batch_ = i_batch]() {
             if constexpr(kHasDropout)
             {
                 return BlockDropout{i_batch_,
                                     i_nhead_,
                                     kargs.num_head_q,
                                     kargs.is_drop_seed_offset_from_host ? kargs.drop_seed.val
                                                                         : *kargs.drop_seed.ptr,
                                     kargs.is_drop_seed_offset_from_host ? kargs.drop_offset.val
                                                                         : *kargs.drop_offset.ptr,
                                     kargs.rp_undrop,
                                     kargs.p_undrop_in_uint8_t,
                                     kargs.is_store_randval};
             }
             else
             {
                 return NullBlockDropout{};
             };
         }();
  
         auto randval_dram_window = [&, i_nhead_ = i_nhead]() {
             constexpr auto randval_dram_window_lengths =
                 make_tuple(number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kN0>{});
             if constexpr(kHasDropout)
             {
                 RandValOutputDataType* rand_val_ptr =
                     reinterpret_cast<RandValOutputDataType*>(kargs.rand_val_ptr) +
                     static_cast<long_index_t>(i_nhead_) * kargs.nhead_stride_randval +
                     batch_offset_randval;
  
                 const auto randval_dram = [&]() {
                     const auto randval_dram_naive =
                         make_naive_tensor_view<address_space_enum::global>(
                             rand_val_ptr,
                             make_tuple(kargs.seqlen_q, kargs.seqlen_k),
                             make_tuple(kargs.stride_randval, 1),
                             number<FmhaPipeline::kAlignmentRandVal>{},
                             number<1>{});
  
                     return pad_tensor_view(randval_dram_naive,
                                            randval_dram_window_lengths,
                                            sequence<kPadSeqLenQ, kPadSeqLenK>{});
                 }();
  
                 return make_tile_window(randval_dram, randval_dram_window_lengths, {i_m0, 0});
             }
             else
             {
                 return make_null_tile_window(randval_dram_window_lengths);
             }
         }();
  
         FmhaMask mask = [&]() {
             if constexpr(kHasMask)
                 return ck_tile::make_generic_attention_mask_from_lr_window<FmhaMask>(
                     kargs.window_size_left,
                     kargs.window_size_right,
                     kargs.sink_size,
                     kargs.seqlen_q,
                     kargs.seqlen_k,
                     kargs.mask_type == GenericAttentionMaskEnum::MASK_FROM_TOP_LEFT);
             else
                 return FmhaMask{kargs.seqlen_q, kargs.seqlen_k};
         }();
  
         // WA i_batch capture structure binding before c++20
         auto position_encoding = [&, i_batch_ = i_batch, i_nhead_ = i_nhead]() {
             if constexpr(BiasEnum == BlockAttentionBiasEnum::ALIBI)
             {
                 // data loading, shared by entire wg
                 // TODO: how to use s_read?
                 SaccDataType slope =
                     *(reinterpret_cast<const SaccDataType*>(kargs.alibi_slope_ptr) +
                       i_batch_ * kargs.alibi_slope_stride + i_nhead_);
 #if CK_TILE_FMHA_FWD_FAST_EXP2
                 slope *= ck_tile::log2e_v<>;
 #endif
                 if constexpr(kHasMask)
                 {
                     return make_alibi_from_lr_mask<SaccDataType, true>(slope,
                                                                        kargs.window_size_left,
                                                                        kargs.window_size_right,
                                                                        kargs.seqlen_q,
                                                                        kargs.seqlen_k,
                                                                        kargs.mask_type);
                 }
                 else
                 {
                     return Alibi<SaccDataType, true>{
                         slope, kargs.seqlen_q, kargs.seqlen_k, AlibiMode::FROM_BOTTOM_RIGHT};
                 }
             }
             else
             {
                 return EmptyPositionEncoding<SaccDataType>{};
             }
         }();
  
         AttentionVariant variant;
         const auto variant_params = [&] {
             const float scale_s = [&] {
                 if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::PERTENSOR)
                 {
                     float q_descale = *(reinterpret_cast<const float*>(kargs.q_descale_ptr));
                     float k_descale = *(reinterpret_cast<const float*>(kargs.k_descale_ptr));
  
                     return kargs.scale_s * q_descale * k_descale;
                 }
                 else
                 {
                     return kargs.scale_s;
                 }
             }();
  
             if constexpr(kHasLogitsSoftCap)
             {
                 return ck_tile::LogitsSoftCapParams<FmhaMask, CK_TILE_FMHA_FWD_FAST_EXP2>{
                     mask, scale_s, kargs.logits_soft_cap, kargs.logits_soft_cap_rcp};
             }
             else
             {
                 return ck_tile::StandardAttentionParams<FmhaMask>{mask, scale_s};
             }
         }();
  
         BlockIndices block_indices{i_batch, i_nhead, i_nhead / kargs.nhead_ratio_qk};
  
         const index_t stride_k_for_pipeline =
             kKVMemoryLayout == BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT
                 ? kVectorSize
                 : kargs.stride_k;
         const index_t stride_v_for_pipeline =
             kKVMemoryLayout == BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT
                 ? kargs.hdim_v
                 : kargs.stride_v;
  
         auto o_acc_tile = [&] {
             if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::PERTENSOR)
             {
                 // TODO - move global load of descale to pipeline
                 float v_descale = *(reinterpret_cast<const float*>(kargs.v_descale_ptr));
  
                 float scale_p = ck_tile::type_convert<float>(ck_tile::numeric<PDataType>::max());
                 float scale_o = v_descale / scale_p;
  
                 auto o_acc_element_func = [&]() {
                     if constexpr(std::is_same_v<ODataType, ck_tile::fp8_t>)
                         return ck_tile::composes(ck_tile::saturates<ck_tile::fp8_t>{},
                                                  ck_tile::scales{scale_o});
                     else
                         return ck_tile::scales{scale_o};
                 }();
  
                 return FmhaPipeline{}(q_dram_window,
                                       identity{}, // q_element_func
                                       k_dram_window,
                                       identity{}, // k_element_func
                                       v_dram_window,
                                       identity{}, // v_element_func
                                       bias_dram_window,
                                       identity{}, // bias_element_func
                                       randval_dram_window,
                                       lse_dram_window,
                                       identity{},         // lse_element_func
                                       identity{},         // s_acc_element_func
                                       scales{scale_p},    // p_compute_element_func
                                       o_acc_element_func, // o_acc_element_func
                                       mask,
                                       position_encoding,
                                       variant_params.sm_scale,
                                       variant,
                                       variant_params,
                                       block_indices,
                                       smem_ptr,
                                       page_idx,
                                       stride_k_for_pipeline,
                                       stride_v_for_pipeline,
                                       kargs.batch_stride_k,
                                       kargs.batch_stride_v,
                                       dropout);
             }
             else
             {
                 return FmhaPipeline{}(q_dram_window,
                                       k_dram_window,
                                       v_dram_window,
                                       bias_dram_window,
                                       randval_dram_window,
                                       lse_dram_window,
                                       mask,
                                       position_encoding,
                                       variant_params.sm_scale,
                                       variant,
                                       variant_params,
                                       block_indices,
                                       smem_ptr,
                                       page_idx,
                                       stride_k_for_pipeline,
                                       stride_v_for_pipeline,
                                       kargs.batch_stride_k,
                                       kargs.batch_stride_v,
                                       dropout);
             }
         }();
  
         // O DRAM and O DRAM window
         auto o_dram = [&]() {
             const auto o_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                 o_ptr,
                 make_tuple(kargs.seqlen_q, kargs.hdim_v),
                 make_tuple(kargs.stride_o, 1),
                 number<FmhaPipeline::kAlignmentO>{},
                 number<1>{});
  
             return pad_tensor_view(
                 o_dram_naive,
                 make_tuple(number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kN1>{}),
                 sequence<kPadSeqLenQ, kPadHeadDimV>{});
         }();
  
         auto o_dram_window =
             make_tile_window(o_dram,
                              make_tuple(number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kN1>{}),
                              {i_m0, i_n1});
  
         EpiloguePipeline{}(o_dram_window, o_acc_tile, nullptr);
     }
 };
  
 } // namespace ck_tile