/*
 * Copyright (c) 2020-2025, NVIDIA CORPORATION.  All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "fmhaRunner.h"
#include "tensorrt_llm/common/envUtils.h"
#include "tensorrt_llm/common/mathUtils.h"
#include <cassert>
#include <cstdio>
#include <cstring>
#include <cuda_runtime.h>
#include <iostream>
#include <math.h>
#include <tuple>
#include <vector>

////////////////////////////////////////////////////////////////////////////////////////////////////

namespace tensorrt_llm
{
namespace kernels
{

////////////////////////////////////////////////////////////////////////////////////////////////////

union __half2_uint32_t_union
{
    half2 fp162;
    uint32_t u32;
};

union __float_uint32_t_union
{
    float fp32;
    uint32_t u32;
};

static inline void set_alpha(uint32_t& alpha, float norm, Data_type dtype)
{
    if (dtype == DATA_TYPE_FP16)
    {
        __half2_uint32_t_union temp;
        temp.fp162 = __float2half2_rn(norm);
        alpha = temp.u32;
    }
    else if (dtype == DATA_TYPE_FP32)
    {
        __float_uint32_t_union temp;
        temp.fp32 = norm;
        alpha = temp.u32;
    }
    else if (dtype == DATA_TYPE_INT32)
    {
        int32_t inorm = static_cast<int32_t>(norm);
        alpha = reinterpret_cast<uint32_t const&>(inorm);
    }
    else if (dtype == DATA_TYPE_BF16)
    {
        // TODO HACK!! BF16 Outputs are computed in FP32 for FP8.
        // This is because cublas does not allow current FP32 output.
        alpha = reinterpret_cast<uint32_t const&>(norm);
    }
    else
    {
        assert(false);
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

FusedMHARunnerV2::FusedMHARunnerV2(MHARunnerFixedParams fixedParams)
    : mFixedParams(fixedParams)
{
    TLLM_CHECK_WITH_INFO((mSM == kSM_80 || mSM == kSM_86 || mSM == kSM_89 || mSM == kSM_90 || mSM == kSM_100
                             || mSM == kSM_120 || mSM == kSM_121),
        "Unsupported architecture");
    TLLM_CHECK_WITH_INFO((mFixedParams.dataType == DATA_TYPE_FP16 || mFixedParams.dataType == DATA_TYPE_BF16
                             || mFixedParams.dataType == DATA_TYPE_E4M3),
        "Unsupported data type");
    xmmaKernel = getXMMAKernelsV2(mFixedParams.dataType, mFixedParams.dataTypeOut, mSM);

    if (mFixedParams.headSizeV == 0)
    {
        mFixedParams.headSizeV = mFixedParams.headSize;
    }
    // Get device attributes.
    int device_id;
    cudaGetDevice(&device_id);
    cudaDeviceGetAttribute(&mMultiProcessorCount, cudaDevAttrMultiProcessorCount, device_id);
    cudaDeviceGetAttribute(&mDeviceL2CacheSize, cudaDevAttrL2CacheSize, device_id);
    auto const [free_memory, total_memory] = tensorrt_llm::common::getDeviceMemoryInfo(false);
    mTotalDeviceMemory = total_memory;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

// Shared setup function.
void FusedMHARunnerV2::setupKernelParams(MHARunnerParams runnerParams)
{
    // Reinit kernel params.
    mKernelParams = {};

    // Set the batch size, and sequence length.
    mKernelParams.b = runnerParams.b;
    mKernelParams.s = runnerParams.qSeqLen;
    mKernelParams.sliding_window_size = runnerParams.slidingWindowSize;
    // Set the log chunked attention size if the chunked attention is used.
    if (mLaunchParams.attention_mask_type == ContextAttentionMaskType::SLIDING_OR_CHUNKED_CAUSAL
        && runnerParams.kvSeqLen > runnerParams.chunkedAttentionSize)
    {
        TLLM_CHECK_WITH_INFO((runnerParams.chunkedAttentionSize & (runnerParams.chunkedAttentionSize - 1)) == 0,
            "Chunked attention size should be a power of 2.");
        mKernelParams.log2_chunked_attention_size = std::log2(runnerParams.chunkedAttentionSize);
    }
    // Set the head size and number of heads.
    mKernelParams.d = mFixedParams.headSize;
    mKernelParams.dv = mFixedParams.headSizeV;
    // The number of grouped heads (only used by generation-phase MLA kernels) currently.
    mKernelParams.num_grouped_heads = runnerParams.numGroupedHeads;
    TLLM_CHECK_WITH_INFO(mFixedParams.numQHeads % mFixedParams.numKvHeads == 0,
        "number of Query heads should be multiple of KV heads !");
    mKernelParams.h = mFixedParams.numQHeads;
    mKernelParams.h_kv = mFixedParams.numKvHeads;
    mKernelParams.h_q_per_kv = mFixedParams.numQHeads / mFixedParams.numKvHeads;
    // Are the input sequences padded ?
    mKernelParams.is_s_padded = mFixedParams.isSPadded;

    // [total_q, h, 2] (max/sum)
    mKernelParams.softmax_stats_ptr = runnerParams.softmaxStatsPtr;
    mKernelParams.softmax_stats_stride_in_bytes = sizeof(float) * 2 * mFixedParams.numQHeads;

    if (mFixedParams.attentionInputLayout == AttentionInputLayout::PACKED_QKV)
    {
        // Packed QKV input layout, [B, S, H * D + H_kv * D + H_kv * Dv].
        mKernelParams.qkv_ptr = runnerParams.qkvPtr;
        mKernelParams.q_stride_in_bytes = mKernelParams.k_stride_in_bytes = mKernelParams.v_stride_in_bytes
            = get_size_in_bytes(mFixedParams.numQHeads * mFixedParams.headSize
                    + mFixedParams.numKvHeads * mFixedParams.headSize
                    + mFixedParams.numKvHeads * mFixedParams.headSizeV,
                mFixedParams.dataType);
    }
    else
    {
        // Contiguous Q input layout, [B, S, H, D].
        mKernelParams.q_ptr = runnerParams.qPtr;
        mKernelParams.q_stride_in_bytes
            = get_size_in_bytes(mFixedParams.numQHeads * mFixedParams.headSize, mFixedParams.dataType);

        // Separate q and kv buffers may have different q and kv sequence lengths.
        mKernelParams.cu_kv_seqlens = reinterpret_cast<int const*>(runnerParams.cuKvSeqLenPtr);

        if (mFixedParams.attentionInputLayout == AttentionInputLayout::Q_CONTIGUOUS_KV)
        {
            // Contiguous kv input layout, [B, S, H_kv * D + H_kv * Dv].
            mKernelParams.kv_ptr = runnerParams.kvPtr;
            mKernelParams.k_stride_in_bytes = mKernelParams.v_stride_in_bytes = get_size_in_bytes(
                mFixedParams.numKvHeads * (mFixedParams.headSize + mFixedParams.headSizeV), mFixedParams.dataType);
        }
        else if (mFixedParams.attentionInputLayout == AttentionInputLayout::Q_PAGED_KV)
        {
            // Paged kv cache layout.
            mKernelParams.paged_kv_cache = runnerParams.pagedKvCache.copyKVBlockArrayForContextFMHA();
            mKernelParams.k_stride_in_bytes = get_size_in_bytes(
                runnerParams.pagedKvCache.mTokensPerBlock * mFixedParams.headSize, mFixedParams.dataType);
            // If d == dv, then v_stride_in_bytes == k_stride_in_bytes.
            // For DeepSeek MLA, which is the only case where d != dv, V is padded to the sizeof K.
            // Thus, v_stride_in_bytes always equals to k_stride_in_bytes so far.
            mKernelParams.v_stride_in_bytes = mKernelParams.k_stride_in_bytes;
        }
    }

    mKernelParams.o_ptr = runnerParams.outputPtr;
    // Set the output buffer stride in bytes.
    mKernelParams.o_stride_in_bytes
        = get_size_in_bytes(mFixedParams.numQHeads * mFixedParams.headSizeV, mFixedParams.dataTypeOut);
    // Set the packed_mask_stride_in_bytes.
    if (mFixedParams.attentionMaskType == ContextAttentionMaskType::CUSTOM_MASK)
    {
        // The packed mask col (n) dimension has to be padded to multiple of 256.
        mKernelParams.packed_mask_stride_in_bytes
            = (tensorrt_llm::common::divUp(int64_t(runnerParams.kvSeqLen), int64_t(FLASH_ATTEN_PACKED_MASK_N_ALIGNMENT))
                  * FLASH_ATTEN_PACKED_MASK_N_ALIGNMENT)
            / 8;
    }

    float const inv_sqrt_scale = (1.f / (sqrtf(mFixedParams.headSize) * mFixedParams.qScaling));
    // Note that we apply scales and bias in the order of
    // (bmm1_output * scale_bmm1 + alibi) * scale_after_alibi
    float const scale_after_alibi = mFixedParams.scaleAlibi ? inv_sqrt_scale : 1.0f;
    float scale_bmm1 = mFixedParams.scaleAlibi ? 1.0f : inv_sqrt_scale;
    // Fuse 1.0f / attn_logit_softcapping_scale into scale_bmm1.
    scale_bmm1 = mFixedParams.attnLogitSoftcappingScale != 0.f ? scale_bmm1 / mFixedParams.attnLogitSoftcappingScale
                                                               : scale_bmm1;
    // The softmax output scale (not used).
    float const scale_softmax = 1.f;
    // FP8 FMHA kernels load the scale_bmm2 from the device memory.
    float const scale_bmm2 = 1.f;

    Data_type scale_type = mLaunchParams.force_fp32_acc ? DATA_TYPE_FP32 : mFixedParams.dataType;
    // Use exp2f optimization for warp-specialized ws kernels on Hopper.
    if (mLaunchParams.useBase2ExpTrick)
    {
        // The kernel adopts the log2f optimization.
        constexpr float kLog2e = 1.4426950408889634074; // log_2(e) = M_LOG2E
        set_alpha(mKernelParams.scale_bmm1, scale_bmm1 * float(kLog2e), DATA_TYPE_FP32);
    }
    else
    {
        set_alpha(mKernelParams.scale_bmm1, scale_bmm1, scale_type);
    }
    set_alpha(mKernelParams.scale_softmax, scale_softmax, scale_type);
    // Host scale_bmm2 will not be used.
    set_alpha(mKernelParams.scale_bmm2, scale_bmm2, scale_type);
    // The attention logit softcapping scale after bmm1 (always float32).
    mKernelParams.softcapping_scale_bmm1 = mFixedParams.attnLogitSoftcappingScale;

    // alibi.
    if (mFixedParams.hasAlibi && mSM > kSM_70)
    {
        mKernelParams.has_alibi = true;
        mKernelParams.alibi_params = AlibiParams(
            mFixedParams.numQHeads, runnerParams.kvSeqLen, mFixedParams.tpSize, mFixedParams.tpRank, scale_after_alibi);
    }

    if (mFixedParams.attentionMaskType == ContextAttentionMaskType::CUSTOM_MASK)
    {
        mKernelParams.packed_mask_ptr = runnerParams.packedMaskPtr;
        mKernelParams.cu_mask_rows = reinterpret_cast<int const*>(runnerParams.cuMaskRowsPtr);
    }
    TLLM_CHECK_WITH_INFO(
        runnerParams.attentionSinksPtr == nullptr || mSM == kSM_90, "The attention sinks is only supported on SM90.");
    mKernelParams.attention_sinks_ptr = runnerParams.attentionSinksPtr;
    mKernelParams.cu_q_seqlens = reinterpret_cast<int const*>(runnerParams.cuQSeqLenPtr);
    mKernelParams.tile_id_counter_ptr = reinterpret_cast<uint32_t*>(runnerParams.tileCounterPtr);
    // TRT doesn't support host scales. Use device scales instead.
    // The scaleBmm1Ptr offset.
    // 2 scales prepared for scaleBmm1 in the device memory: float scale, float (scale with log2e).
    int64_t scaleBmm1PtrOffset = (mLaunchParams.useBase2ExpTrick ? kIdxScaleSoftmaxLog2Ptr : kIdxScaleSoftmaxPtr);
    // Only fp8 kernels need to load scales from the device memory.
    if (mFixedParams.dataType == DATA_TYPE_E4M3)
    {
        mKernelParams.scale_bmm1_d = reinterpret_cast<uint32_t const*>(runnerParams.scaleBmm1Ptr + scaleBmm1PtrOffset);
        mKernelParams.scale_bmm2_d = reinterpret_cast<uint32_t const*>(runnerParams.scaleBmm2Ptr);
    }

    // for sage attention
    mKernelParams.sage.q.scales = runnerParams.qScalePtr;
    mKernelParams.sage.k.scales = runnerParams.kScalePtr;
    mKernelParams.sage.v.scales = runnerParams.vScalePtr;
    mKernelParams.sage.q.max_nblock = runnerParams.qMaxNBlock;
    mKernelParams.sage.k.max_nblock = runnerParams.kMaxNBlock;
    mKernelParams.sage.v.max_nblock = runnerParams.vMaxNBlock;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

// Set the launch params to select kernels.
void FusedMHARunnerV2::setupLaunchParams(MHARunnerParams runnerParams)
{

    // Determine launch parameters.
    // Reset launch params to default.
    mLaunchParams = {};

    // Device properties.
    mLaunchParams.multi_processor_count = mMultiProcessorCount;
    mLaunchParams.device_l2_cache_size = mDeviceL2CacheSize;
    mLaunchParams.total_device_memory = mTotalDeviceMemory;

    // Do we use attnLogitSoftcappingScale ?
    TLLM_CHECK_WITH_INFO(
        (mFixedParams.headSize == 128 || mFixedParams.headSize == 256) || !mFixedParams.attnLogitSoftcappingScale,
        "FMHA only supports head_size = 128 or 256 with attention logit softcapping scale currently.");
    mLaunchParams.enableAttnLogitSoftcapping = mFixedParams.attnLogitSoftcappingScale != 0.f;
    // BF16 FMHA only accumulates on FP32.
    // E4M3 FMHA only supports fp32 accumulation currently.
    mLaunchParams.force_fp32_acc = mFixedParams.dataType == DATA_TYPE_BF16 || mFixedParams.dataType == DATA_TYPE_E4M3
        || mFixedParams.forceFp32Acc || runnerParams.forceFp32Acc;
    // The attention mask type.
    mLaunchParams.attention_mask_type = mFixedParams.attentionMaskType;
    // The input layout type.
    mLaunchParams.attention_input_layout = mFixedParams.attentionInputLayout;

    // The total sequence length used to set the tma descriptors.
    mLaunchParams.total_q_seqlen
        = mFixedParams.isSPadded ? runnerParams.b * runnerParams.qSeqLen : runnerParams.totalQSeqLen;
    mLaunchParams.total_kv_seqlen
        = mFixedParams.isSPadded ? runnerParams.b * runnerParams.kvSeqLen : runnerParams.totalKvSeqLen;
    // Workaround for nvbug 5412456: total_kv_seqlen fallbacks to total_q_seqlen if it's zero.
    if (mLaunchParams.total_kv_seqlen == 0)
    {
        mLaunchParams.total_kv_seqlen = mLaunchParams.total_q_seqlen;
    }

    TLLM_CHECK_WITH_INFO(mFixedParams.headSize > 0, "Head size should be greater than 0.");
    // Pad head size to next power of 2.
    int padded_d_next_power_of_2 = (mFixedParams.headSize & (mFixedParams.headSize - 1)) == 0
        ? mFixedParams.headSize
        : pow(2, int(log2(mFixedParams.headSize)) + 1);
    // In fact, due to 128B swizzle mode of TMA, only 128 bytes alignment is required,
    // so we pad head size to next multiply of 128B.
    int d_per_group = 128 / get_size_in_bytes(mFixedParams.dataType);
    int d_groups = (mFixedParams.headSize + d_per_group - 1) / d_per_group;
    int padded_d_next_multiply_of_128byte = d_groups * d_per_group;
    // Choose the smaller one to save SMEM.
    mLaunchParams.padded_d = std::min(padded_d_next_power_of_2, padded_d_next_multiply_of_128byte);

    bool const isSm70 = (mSM == kSM_70);
    bool const isSm90 = (mSM == kSM_90);
    bool const isSm8x = (mSM == kSM_86 || mSM == kSM_89);
    bool const isSm80 = (mSM == kSM_80);
    bool const isSm89 = (mSM == kSM_89);
    bool const isSm100 = (mSM == kSM_100);
    bool const isSm120f = (mSM == kSM_120 || mSM == kSM_121);

    // Sliding_or_chunked_causal mask.
    if ((runnerParams.kvSeqLen > runnerParams.slidingWindowSize
            || runnerParams.kvSeqLen > runnerParams.chunkedAttentionSize)
        && mLaunchParams.attention_mask_type == ContextAttentionMaskType::CAUSAL)
    {
        TLLM_CHECK_WITH_INFO(!(runnerParams.kvSeqLen > runnerParams.chunkedAttentionSize
                                 && runnerParams.kvSeqLen > runnerParams.slidingWindowSize),
            "Chunked attention size and sliding window size should not be used together.");
        TLLM_CHECK_WITH_INFO(isSm90 || runnerParams.kvSeqLen <= runnerParams.chunkedAttentionSize,
            "Chunked attention is only supported on Sm90.");
        mLaunchParams.attention_mask_type = ContextAttentionMaskType::SLIDING_OR_CHUNKED_CAUSAL;
    }

    // Is the input layout separate q + kv input ?
    bool const separateQKvInput = mFixedParams.attentionInputLayout != AttentionInputLayout::PACKED_QKV;
    // Is the mask type padding or causal mask ?
    bool const paddingOrCausalMask = mFixedParams.attentionMaskType == ContextAttentionMaskType::PADDING
        || mFixedParams.attentionMaskType == ContextAttentionMaskType::CAUSAL;

    // Only warp-specialized FMHA kernels support FP8 on Hopper.
    // Separate Q + KV input layout: enable warp-specialization kernels when s > 512, otherwise use ampere-style flash
    // attention kernels.
    if (isSm90 && (mFixedParams.dataType == DATA_TYPE_E4M3 || (separateQKvInput && runnerParams.kvSeqLen > 512)))
    {
        mLaunchParams.flash_attention = true;
        mLaunchParams.force_unroll = true;
    }
    else if (isSm70)
    {
        TLLM_CHECK_WITH_INFO(false, "Unsupported architecture");
    }
    // Hopper: fallback to original fmha_v2 when head_size <= 64 and seq_len <= 256
    // Only supports packed_qkv input + padding/causal mask.
    else if (isSm90 && !separateQKvInput && paddingOrCausalMask
        && (mFixedParams.headSize == 32 || mFixedParams.headSize == 64) && runnerParams.qSeqLen <= 256
        && !common::getEnvForceDeterministicAttention())
    {
        mLaunchParams.flash_attention = false;
        // get max sequence length for non-flash-attention.
        // this doesn't support different q and kv sequence lengths.
        mLaunchParams.kernel_s = getSFromMaxSeqLen(runnerParams.qSeqLen);
    }
    else
    { // always use flash attention kernels for Ampere/Ada
        mLaunchParams.flash_attention = true;
        // flash attention kernles s = 0 (support any seq length)
        mLaunchParams.kernel_s = 0;
        mLaunchParams.force_unroll = true;
        // enable tiled kernels on Ampere/Ada
        if ((isSm89 || isSm120f) && mFixedParams.dataType == DATA_TYPE_E4M3)
        {
            // so far Ada QMMA only supports non-tiled kernels.
            mLaunchParams.granular_tiling = false;
        }
        else if (mLaunchParams.flash_attention && runnerParams.kvSeqLen <= 64)
        {
            // flash attention tiled kernels allows larger free dim tile size (M, N) with flexibility
            // in unroll dimension tile size (K). for short sequence length (s<=128), tiled kernels
            // can suffer from tile quantization loss therefore use flash attention non-tiled instead
            mLaunchParams.granular_tiling = false;
        }
        else if ((isSm8x || isSm120f) && mFixedParams.headSize < 256)
        {
            // flash attention tiled kernel is faster on Ada and Ampere derivatives when head_size>=256
            mLaunchParams.granular_tiling = false;
        }
        else if (isSm80 || isSm8x || isSm100 || isSm120f)
        {
            // otherwise, choose tiled kernel for Ampere/Ada/Gb20x
            mLaunchParams.granular_tiling = true;
        }
    }

    // when flash attention is enabled on Hopper, we need to set the tma descriptors
    if (isSm90 && mLaunchParams.flash_attention)
    {
        mLaunchParams.warp_specialization = true;
        mLaunchParams.use_tma = true;
        // Enable dynamic tile scheduling for hopper ws kernel.
        mLaunchParams.dynamic_scheduler = true;
    }

    // Use specialized ws kernels on Hopper for cases without alibi.
    if (mLaunchParams.warp_specialization && !mFixedParams.hasAlibi)
    {
        // Use specialized ws kernels for cases without alibi.
        mLaunchParams.useKernelWithoutAlibi = true;
        // Enable exp2f optimization (which helps improve performance).
        //    - note that this is not compatible with alibi bias due to the accuracy issues.
        //    - only hopper warp-specialized kernels have this optimization.
        //    - it doesn't work with attention logit softcapping.
        mLaunchParams.useBase2ExpTrick = !mLaunchParams.enableAttnLogitSoftcapping;
    }

    // TODO: Refactor these dirty hacks.
    // For Deepseek-v2(MLA), all of SM80, SM89 and SM90 kernels use tiled flash attention
    // in both context (192/128 dimensions) and generation (576/512 dimensions)
    if (mFixedParams.headSize == mFixedParams.headSizeV + 64)
    {
        mLaunchParams.flash_attention = true;
        mLaunchParams.force_unroll = true;
        mLaunchParams.kernel_s = 0;

        // Now we have SM90 FP8 generation and BF16 context MLA kernels
        bool isHopperFP8GenerationMLA
            = isSm90 && mFixedParams.dataType == DATA_TYPE_E4M3 && mFixedParams.headSizeV == 512;
        bool isHopperBF16ContextMLA
            = isSm90 && mFixedParams.dataType == DATA_TYPE_BF16 && mFixedParams.headSizeV == 128;

        // These treatments are only for other MLA cases
        if (!isHopperFP8GenerationMLA && !isHopperBF16ContextMLA)
        {
            mLaunchParams.granular_tiling = true;
            // Even on SM90, we use ampere-style kernel, will be optimized later
            mLaunchParams.warp_specialization = false;
            mLaunchParams.useKernelWithoutAlibi = false;
            // Deepseek-V2 kernel is not hooper style right now.
            mLaunchParams.useBase2ExpTrick = false;
            mLaunchParams.use_tma = false;
            mLaunchParams.dynamic_scheduler = false;
        }
    }

    mLaunchParams.sage_block_size_q = mFixedParams.sageBlockSizeQ;
    mLaunchParams.sage_block_size_k = mFixedParams.sageBlockSizeK;
    mLaunchParams.sage_block_size_v = mFixedParams.sageBlockSizeV;
    // for not (sm90 + warp_specialization + flash attention kernel) kernel:
    //   all kernels enable saving softmaxStatsPtr, just let softmaxStatsPtr != null
    // for (sm90 + warp_specialization + flash attention) kernel:
    //   we need to explicitly set supportReturnSoftmaxStats to true when
    //  satisfying the following constrains
    if (!isSm90)
    {
        mLaunchParams.supportReturnSoftmaxStats = true;
    }
    else
    {
        bool isHopperBF16ContextMLA = (mFixedParams.headSize == mFixedParams.headSizeV + 64) && isSm90
            && mFixedParams.dataType == DATA_TYPE_BF16 && mFixedParams.headSizeV == 128;
        // TODO: add support for separate QKV input layout
        mLaunchParams.supportReturnSoftmaxStats = (runnerParams.softmaxStatsPtr != nullptr
            && mLaunchParams.flash_attention && mLaunchParams.warp_specialization
            && ((!isHopperBF16ContextMLA
                    && mLaunchParams.attention_input_layout == AttentionInputLayout::Q_CONTIGUOUS_KV)
                || (isHopperBF16ContextMLA
                    && (mLaunchParams.attention_input_layout == AttentionInputLayout::Q_PAGED_KV))));
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

// TMA descriptors are used as grid_constant parameters (remove MemCpyH2D operations)
void FusedMHARunnerV2::setTmaDescriptors(MHARunnerParams runnerParams)
{
    const uint32_t d = mKernelParams.d;
    const uint32_t dv = mKernelParams.dv;
    const uint32_t h = mKernelParams.h;
    const uint32_t h_kv = mKernelParams.h_kv;
    const uint32_t total_q_seqlen = mLaunchParams.total_q_seqlen;
    const uint32_t total_kv_seqlen = mLaunchParams.total_kv_seqlen;

    uint64_t const d_in_bytes = get_size_in_bytes(d, mFixedParams.dataType);
    uint64_t const dv_in_bytes = get_size_in_bytes(dv, mFixedParams.dataType);

    // split D into multiple groups in order to match the TMA swizzle mode (128B)
    uint32_t const padded_d_in_bytes = get_size_in_bytes(mLaunchParams.padded_d, mFixedParams.dataType);
    uint32_t const d_groups = padded_d_in_bytes > 128 ? padded_d_in_bytes / 128 : 1;
    uint32_t const d_bytes_per_group = padded_d_in_bytes / d_groups;
    uint32_t const d_per_group = mLaunchParams.padded_d / d_groups;

    uint32_t q_step = 0, kv_step = 0;
    xmmaKernel->getStepSize(q_step, kv_step, mKernelParams, mLaunchParams);

    auto const layout = mFixedParams.attentionInputLayout;

    // Q Layout: [total_seqlen, H, D]
    const uint32_t tensor_size_q[3] = {d, h, total_q_seqlen};

    // Stride size in bytes. Assumes least significant dim is 1
    const uint64_t tensor_stride_q[2] = {d_in_bytes, uint64_t(mKernelParams.q_stride_in_bytes)};

    // Starting memory address
    char const* q_ptr = reinterpret_cast<char const*>(
        layout == AttentionInputLayout::PACKED_QKV ? mKernelParams.qkv_ptr : mKernelParams.q_ptr);

    // Box size of TMA
    const uint32_t box_size_q[3] = {d_per_group, 1, q_step};

    // Traversal stride.
    const uint32_t traversal_stride[3] = {1, 1, 1};

    // OOB fill zeros.
    const uint32_t oob_fill = 0;

    // FP32 to TF32 conversion disabled.
    const uint32_t fp32_to_tf32 = 0;

    // GMMA descriptor mode.
    cudaTmaDescSwizzle const swizzle_mode = (d_bytes_per_group > 64
            ? cudaTmaDescSwizzle::SWIZZLE_128B
            : (d_bytes_per_group > 32 ? cudaTmaDescSwizzle::SWIZZLE_64B : cudaTmaDescSwizzle::SWIZZLE_32B));

    // Desc Format (data type).
    cudaTmaDescFormat const desc_format
        = (get_size_in_bytes(mFixedParams.dataType) == 1) ? cudaTmaDescFormat::U8 : cudaTmaDescFormat::F16_RN;

    Multiple_tma_descriptor<3> qo_tma_descriptor;

    // Q
    qo_tma_descriptor.set_tma_desctriptor(q_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode,
        cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_q, tensor_stride_q, traversal_stride, box_size_q,
        oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_q);

    // O
    if ((get_size_in_bytes(mFixedParams.dataTypeOut) == 1)
        && mLaunchParams.attention_mask_type != ContextAttentionMaskType::SLIDING_OR_CHUNKED_CAUSAL)
    {
        // O Layout: [total_seqlen, H, DV]
        const uint32_t tensor_size_o[3] = {dv, h, total_q_seqlen};

        const uint64_t tensor_stride_o[2]
            = {get_size_in_bytes(dv, mFixedParams.dataTypeOut), uint64_t(mKernelParams.o_stride_in_bytes)};

        char* o_ptr = reinterpret_cast<char*>(mKernelParams.o_ptr);

        // Box size of TMA
        const uint32_t box_size_o[3] = {d_per_group, 1, 16};

        // dataTypeOut may be different with dataType, so desc_format and swizzle_mode
        // may be incorrect. For example, QKV are in bf16 while O is in fp8.
        // Luckily, this case doesn't exist so far. But we should keep one eye on it.
        qo_tma_descriptor.set_tma_desctriptor(o_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED,
            swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_o, tensor_stride_o, traversal_stride,
            box_size_o, oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_o);
    }

    if (layout == AttentionInputLayout::Q_PAGED_KV)
    {
        // KV in q_paged_kv uses 4D tensor
        // Layout: [INT32_MAX, H_KV, TokensPerBlock, D]
        const uint32_t tokens_per_block = mKernelParams.paged_kv_cache.mTokensPerBlock;
        const uint32_t tensor_size_k[4] = {d, tokens_per_block, h_kv, INT_MAX};
        const uint32_t tensor_size_v[4] = {dv, tokens_per_block, h_kv, INT_MAX};

        const uint64_t tensor_stride_k[3] = {uint64_t(mKernelParams.k_stride_in_bytes / tokens_per_block), // d
            uint64_t(mKernelParams.k_stride_in_bytes),                                                     // d * 64
            uint64_t(mKernelParams.paged_kv_cache.mBytesPerBlock)};
        const uint64_t tensor_stride_v[3]
            = {// we cannot use dv * Kernel_traits::ELEMENT_BYTES because V may be padded (MLA)
                uint64_t(mKernelParams.v_stride_in_bytes / tokens_per_block), // dv
                uint64_t(mKernelParams.v_stride_in_bytes),                    // dv * 64
                uint64_t(mKernelParams.paged_kv_cache.mBytesPerBlock)};

        char const* kv_ptr = reinterpret_cast<char*>(runnerParams.pagedKvCache.mPrimaryPoolPtr);

        const uint32_t box_size_kv[4] = {d_per_group, std::min(tokens_per_block, kv_step), 1, 1};

        TLLM_CHECK(kv_step % tokens_per_block == 0 || tokens_per_block % kv_step == 0);
        mKernelParams.blocks_per_tma_load = std::max<uint32_t>(1, kv_step / tokens_per_block);
        mKernelParams.blocks_per_tma_load_log2 = log2(mKernelParams.blocks_per_tma_load);

        const uint32_t traversal_stride[4] = {1, 1, 1, 1};

        Multiple_tma_descriptor<4> kv_tma_descriptor;
        // K
        kv_tma_descriptor.set_tma_desctriptor(kv_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED,
            swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_k, tensor_stride_k, traversal_stride,
            box_size_kv, oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_k);
        // V
        kv_tma_descriptor.set_tma_desctriptor(kv_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED,
            swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_v, tensor_stride_v, traversal_stride,
            box_size_kv, oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_v);
    }
    else
    {
        // Otherwise KV uses 3D tensor
        const uint32_t tensor_size_k[3] = {d, h_kv, total_kv_seqlen};
        const uint32_t tensor_size_v[3] = {dv, h_kv, total_kv_seqlen};

        const uint64_t tensor_stride_k[2] = {d_in_bytes, uint64_t(mKernelParams.k_stride_in_bytes)};
        const uint64_t tensor_stride_v[2] = {dv_in_bytes, uint64_t(mKernelParams.v_stride_in_bytes)};

        const uint32_t box_size_kv[3] = {d_per_group, 1, kv_step};

        char const *k_ptr, *v_ptr;

        if (layout == AttentionInputLayout::PACKED_QKV)
        {
            // Layout: [total_seqlen, (H, D) + (H_KV, D) + (H_KV, DV)]
            k_ptr = q_ptr + h * d_in_bytes;
            v_ptr = k_ptr + h_kv * d_in_bytes;
        }
        else if (layout == AttentionInputLayout::Q_CONTIGUOUS_KV)
        {
            // Layout, [B, S, H_kv * D + H_kv * Dv].
            k_ptr = reinterpret_cast<char const*>(mKernelParams.kv_ptr);
            v_ptr = k_ptr + h_kv * d_in_bytes;
        }

        Multiple_tma_descriptor<3> kv_tma_descriptor;
        // K
        kv_tma_descriptor.set_tma_desctriptor(k_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED,
            swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_k, tensor_stride_k, traversal_stride,
            box_size_kv, oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_k);
        // V
        kv_tma_descriptor.set_tma_desctriptor(v_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED,
            swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_v, tensor_stride_v, traversal_stride,
            box_size_kv, oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_v);
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////
void FusedMHARunnerV2::run(MHARunnerParams runnerParams)
{
    // Note that we must set the launch params first.
    // Set the launch params.
    setupLaunchParams(runnerParams);
    // Set the kernel params.
    setupKernelParams(runnerParams);
    // Need to set tma descriptors additionally.
    if (mSM == kSM_90 && mLaunchParams.use_tma)
    {
        setTmaDescriptors(runnerParams);
    }
    // Select the kernel and run it.
    xmmaKernel->run(mKernelParams, mLaunchParams, runnerParams.stream);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

bool FusedMHARunnerV2::isValidS(int s) const
{
    return xmmaKernel->isValid(s);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

int FusedMHARunnerV2::getSFromMaxSeqLen(int const max_seq_len) const
{
    int S = 1024;

    if (max_seq_len <= 64)
    {
        S = 64;
    }
    else if (max_seq_len <= 128)
    {
        S = 128;
    }
    else if (max_seq_len <= 256)
    {
        S = 256;
    }
    else if (max_seq_len <= 384)
    {
        S = 384;
    }
    else if (max_seq_len <= 512)
    {
        S = 512;
    }
    // for bert and vit, use flash attention when s >= 512
    else if (max_seq_len > 512)
    {
        S = max_seq_len;
    }

    return S;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

// Function to check if fmha is supported when building plugins.
// If any kernel in the map meets the requirements, then return true.
bool FusedMHARunnerV2::isFmhaSupported()
{
    bool is_supported = xmmaKernel->checkIfKernelExist(mFixedParams);
    if (!is_supported)
    {
        std::string msg = "FMHA Kernel doesn't exist for mFixedParams:\n" + mFixedParams.convertToStrOutput();
        TLLM_LOG_WARNING("%s\n", msg.c_str());
    }
    return is_supported;
}

} // namespace kernels
} // namespace tensorrt_llm
