{{ copyright }}

#include <fused_multihead_attention_utils.h>
#include <fmha/hopper/gmma_descriptor.h>
#include <fmha/hopper/smem_tile.h>
#include <fmha/utils.h>
#include <fmha/hopper/compute_tile.h>

#include <fmha/warpspec/kernel_traits.h>
#include <fmha/warpspec/dma.h>
#include <fmha/warpspec/compute.h>

{{ include_str }}
////////////////////////////////////////////////////////////////////////////////////////////////////
{{ local_ns_open }}
#if CUDA_VERSION >= {{ min_cuda_version }}

static constexpr int DMA2COMPUTE_DEPTH = 1;
{{ num_compute_groups_str }}
static constexpr bool USE_TMA_STORE = {{ use_tma_store_flag }};

{{ bert_launch_params }}
{{ attn_mask_type_str }}

using Ktraits = {{ kernel_traits_header }}
                {{ loop_step }},
                {{ kv_loop_step }},
                {{ head_size }},
                {{ head_size_v }},
                {{ q_tile_buffers }},
                {{ kv_tile_buffers }},
                NUM_COMPUTE_GROUPS,
                DMA2COMPUTE_DEPTH,
                0,
                {{ heads_interleaved_flag }},
                false,
                {{ enable_mutex_flag }},
                {{ scheduling_mode }},
                {{ input_layout_flag }},
                USE_TMA_STORE,
                {{ enable_attn_logit_softcapping_flag }},
                {{ return_softmax_stats_flag }},
                {{ enable_skip_softmax_flag }},
                {{ output_dtype_ }},
                {{ sage_block_size_q }},
                {{ sage_block_size_k }},
                {{ sage_block_size_v }}>;

using Ktraits_causal = {{ kernel_traits_header }}
                       {{ loop_step }},
                       {{ kv_loop_step }},
                       {{ head_size }},
                       {{ head_size_v }},
                       {{ q_tile_buffers }},
                       {{ kv_tile_buffers }},
                       NUM_COMPUTE_GROUPS,
                       DMA2COMPUTE_DEPTH,
                       1,
                       {{ heads_interleaved_flag }},
                       {{ has_alibi }},
                       {{ enable_mutex_flag }},
                       {{ scheduling_mode }},
                       {{ input_layout_flag }},
                       USE_TMA_STORE,
                       {{ enable_attn_logit_softcapping_flag }},
                       {{ return_softmax_stats_flag }},
                       {{ enable_skip_softmax_flag }},
                       {{ output_dtype_ }}>;

using Ktraits_sliding_or_chunked_causal = {{ kernel_traits_header }}
                                      {{ loop_step }},
                                      {{ kv_loop_step }},
                                      {{ head_size }},
                                      {{ head_size_v }},
                                      {{ q_tile_buffers }},
                                      {{ kv_tile_buffers }},
                                      NUM_COMPUTE_GROUPS,
                                      DMA2COMPUTE_DEPTH,
                                      2,
                                      {{ heads_interleaved_flag }},
                                      {{ has_alibi }},
                                      {{ enable_mutex_flag }},
                                      {{ scheduling_mode }},
                                      {{ input_layout_flag }},
                                      USE_TMA_STORE && false,
                                      {{ enable_attn_logit_softcapping_flag }},
                                      {{ return_softmax_stats_flag }},
                                      {{ enable_skip_softmax_flag }},
                                      {{ output_dtype_ }}>;

using Ktraits_custom_mask = {{ kernel_traits_header }}
                            {{ loop_step }},
                            {{ kv_loop_step }},
                            {{ head_size }},
                            {{ head_size_v }},
                            {{ q_tile_buffers }},
                            {{ kv_tile_buffers }},
                            NUM_COMPUTE_GROUPS,
                            DMA2COMPUTE_DEPTH,
                            3,
                            {{ heads_interleaved_flag }},
                            {{ has_alibi }},
                            {{ enable_mutex_flag }},
                            {{ scheduling_mode }},
                            {{ input_layout_flag }},
                            USE_TMA_STORE && false,
                            {{ enable_attn_logit_softcapping_flag }},
                            {{ return_softmax_stats_flag }},
                            {{ enable_skip_softmax_flag }},
                            {{ output_dtype_ }}>;

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

{% if padding_mask %} // padding_mask

using Shared = typename Ktraits::Shared;

extern "C"
__global__ __launch_bounds__(Ktraits::THREADS, 1)
void {{ kernel_name }}(
    const __grid_constant__ {{ params_type }} params){

    extern __shared__ char smem_[];
    char *smem_aligned = fmha::align_1024(smem_);

    Shared *shared = reinterpret_cast<Shared *>(&smem_aligned[0]);
    shared->init(threadIdx.x == 0);
    __syncthreads();

    // special trick to avoid wrap_sync (leads to illegal instruction)
    int warp_group = __shfl_sync(0xffffffff, threadIdx.x / 128, 0);
    int tidx = threadIdx.x % 128;

    if( warp_group == NUM_COMPUTE_GROUPS ) {  // dma + sched

        {{ setmaxnreg_dma_str }}
        uint32_t elect_one = tidx == 0;

        // Need all threads involved when the dam group needs to transpose the v tile explicltly.
        if constexpr ( Ktraits::DMA_GROUP_TRANSPOSE_V ) {
            fmha::ws::DMA<Ktraits>::Device dma_device(elect_one);
            dma_device.{{ run_fct_name }}(params, shared);
        } else {
            fmha::ws::DMA<Ktraits>::Device dma_device(elect_one);
            if( tidx < 32 ) {
                dma_device.{{ run_fct_name }}(params, shared);
            }
        }

    } else {  // math

        {{ setmaxnreg_compute_str }}

        fmha::ws::Compute<fmha::{{ instruction_traits }}, Ktraits> compute;
        compute.run(warp_group, tidx, shared, params);
    }
}

{% endif %} // padding mask

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

{% if causal_mask %} // causal_mask

using Shared_causal = typename Ktraits_causal::Shared;

extern "C"
__global__ __launch_bounds__(Ktraits_causal::THREADS, 1)
void {{ causal_kernel_name }}(
    const __grid_constant__ {{ params_type }} params){

    extern __shared__ char smem_[];
    char *smem_aligned = fmha::align_1024(smem_);

    Shared_causal *shared = reinterpret_cast<Shared_causal *>(&smem_aligned[0]);
    shared->init(threadIdx.x == 0);
    __syncthreads();

    // special trick to avoid wrap_sync (leads to illegal instruction)
    int warp_group = __shfl_sync(0xffffffff, threadIdx.x / 128, 0);
    int tidx = threadIdx.x % 128;

    if( warp_group == NUM_COMPUTE_GROUPS ) {  // dma + sched

        {{ setmaxnreg_dma_str }}
        uint32_t elect_one = tidx == 0;

        // Need all threads involved when the dam group needs to transpose the v tile explicltly.
        if constexpr ( Ktraits_causal::DMA_GROUP_TRANSPOSE_V ) {
            fmha::ws::DMA<Ktraits_causal>::Device dma_device(elect_one);
            dma_device.{{ run_fct_name }}(params, shared);
        } else {
            fmha::ws::DMA<Ktraits_causal>::Device dma_device(elect_one);
            if( tidx < 32 ) {
                dma_device.{{ run_fct_name }}(params, shared);
            }
        }

    } else {  // math

        {{ setmaxnreg_compute_str }}

        fmha::ws::Compute<fmha::{{ instruction_traits }}, Ktraits_causal> compute;
        compute.run(warp_group, tidx, shared, params);
    }
}

{% endif %} // causal mask

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

{% if sliding_or_chunked_causal_mask %} // sliding_or_chunked_causal_mask

using Shared_sliding_or_chunked_causal = typename Ktraits_sliding_or_chunked_causal::Shared;

extern "C"
__global__ __launch_bounds__(Ktraits_sliding_or_chunked_causal::THREADS, 1)
void {{ sliding_or_chunked_causal_kernel_name }}(
    const __grid_constant__ {{ params_type }} params){

    extern __shared__ char smem_[];
    char *smem_aligned = fmha::align_1024(smem_);

    Shared_sliding_or_chunked_causal *shared =
        reinterpret_cast<Shared_sliding_or_chunked_causal *>(&smem_aligned[0]);
    shared->init(threadIdx.x == 0);
    __syncthreads();

    // special trick to avoid wrap_sync (leads to illegal instruction)
    int warp_group = __shfl_sync(0xffffffff, threadIdx.x / 128, 0);
    int tidx = threadIdx.x % 128;

    if( warp_group == NUM_COMPUTE_GROUPS ) {  // dma + sched

        {{ setmaxnreg_dma_str }}
        uint32_t elect_one = tidx == 0;

        // Need all threads involved when the dam group needs to transpose the v tile explicltly.
        if constexpr ( Ktraits_sliding_or_chunked_causal::DMA_GROUP_TRANSPOSE_V ) {
            fmha::ws::DMA<Ktraits_sliding_or_chunked_causal>::Device dma_device(elect_one);
            dma_device.{{ run_fct_name }}(params, shared);
        } else {
            fmha::ws::DMA<Ktraits_sliding_or_chunked_causal>::Device dma_device(elect_one);
            if( tidx < 32 ) {
                dma_device.{{ run_fct_name }}(params, shared);
            }
        }

    } else {  // math

        {{ setmaxnreg_compute_str }}

        fmha::ws::Compute<fmha::{{ instruction_traits }}, Ktraits_sliding_or_chunked_causal> compute;
        compute.run(warp_group, tidx, shared, params);
    }
}

{% endif %} // sliding_or_chunked_causal_mask

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

{% if custom_mask %} // custom_mask

using Shared_custom_mask = typename Ktraits_custom_mask::Shared;

extern "C"
__global__ __launch_bounds__(Ktraits_custom_mask::THREADS, 1)
void {{ custom_mask_kernel_name }}(
    const __grid_constant__ {{ params_type }} params){

    extern __shared__ char smem_[];
    char *smem_aligned = fmha::align_1024(smem_);

    Shared_custom_mask *shared =
        reinterpret_cast<Shared_custom_mask *>(&smem_aligned[0]);
    shared->init(threadIdx.x == 0);
    __syncthreads();

    // special trick to avoid wrap_sync (leads to illegal instruction)
    int warp_group = __shfl_sync(0xffffffff, threadIdx.x / 128, 0);
    int tidx = threadIdx.x % 128;

    if( warp_group == NUM_COMPUTE_GROUPS ) {  // dma + sched

        {{ setmaxnreg_dma_str }}
        uint32_t elect_one = tidx == 0;

        // Need all threads involved when the dam group needs to transpose the v tile explicltly.
        if constexpr ( Ktraits_custom_mask::DMA_GROUP_TRANSPOSE_V ) {
            fmha::ws::DMA<Ktraits_custom_mask>::Device dma_device(elect_one);
            dma_device.{{ run_fct_name }}(params, shared);
        } else {
            fmha::ws::DMA<Ktraits_custom_mask>::Device dma_device(elect_one);
            if( tidx < 32 ) {
                dma_device.{{ run_fct_name }}(params, shared);
            }
        }

    } else {  // math

        {{ setmaxnreg_compute_str }}

        fmha::ws::Compute<fmha::{{ instruction_traits }}, Ktraits_custom_mask> compute;
        compute.run(warp_group, tidx, shared, params);
    }
}

{% endif %} // custom_mask

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

void {{ launcher_name }}(
    {{ fused_multihead_attention_params_v2_str }} &params,
    const Launch_params &launch_params, cudaStream_t stream){

    {{ TMA_config }}
    if( Ktraits::SCHEDULING_MODE > 0 ) {
        FMHA_CHECK_CUDA(cudaMemsetAsync(params.tile_id_counter_ptr, 0, sizeof(uint32_t), stream));
    }

    dim3 block_size;

    if( Ktraits::SCHEDULING_MODE == 0 ) {
        block_size.y = std::min(params.b * params.h, launch_params.multi_processor_count);
        // distribute m steps to multiple blocks (fully utilize SMs)
        // block.x = blocks that handle single head, block.y = blocks that handle different heads
        size_t sms_per_head = (launch_params.multi_processor_count) / block_size.y;
        // Take multiple compute groups into consideration.
        size_t m_steps = size_t((params.s + {{ loop_step }} * NUM_COMPUTE_GROUPS - 1) / ({{ loop_step }} * NUM_COMPUTE_GROUPS));

        // 2 * {{ bytes_per_elt }} stands for kv cache and {{ bytes_per_elt }} bytes per element.
        size_t size_in_bytes = static_cast<size_t>(block_size.y) * params.s * params.d * 2 * {{ bytes_per_elt }};
        if( size_in_bytes <= launch_params.device_l2_cache_size ) {
            // strategy 1: limit to only 1 wave
            block_size.x = std::min(m_steps, sms_per_head);
        } else {
            // strategy 2: fully unroll the q loops (contiguous blocks handle all q loops)
            block_size.x = m_steps;
        }
        params.num_tiles = params.b * params.h;
    } else if( Ktraits::SCHEDULING_MODE == 1 ) {
        // Get the max total M steps
        // Take multiple compute groups into consideration.
        size_t m_steps = size_t((params.s + {{ loop_step }} * NUM_COMPUTE_GROUPS - 1) / ({{ loop_step }} * NUM_COMPUTE_GROUPS));
        params.num_tiles_per_head = static_cast<uint32_t>(m_steps);
        params.num_tiles = static_cast<uint32_t>(m_steps * params.b * params.h);
        if (launch_params.attention_mask_type == Attention_mask_type::CAUSAL) {
            // 2 * {{ bytes_per_elt }} stands for kv cache and {{ bytes_per_elt }} bytes per element.
            size_t size_in_bytes = static_cast<size_t>(params.b) * params.h * params.s * params.d * 2 * {{ bytes_per_elt }};
            params.use_balanced_scheduling = {{ "true" if fp8_kernel else "(size_in_bytes <= launch_params.device_l2_cache_size)" }};
            // Head-first tile order streams many heads' KV at once, so require the whole KV
            // working set to fit in L2 (with 2x headroom, a measured heuristic).
            // h_kv counts each GQA group's shared KV once. total_kv_seqlen is
            // b * max_kv_len, so past KV (chunked prefill) is measured; fall back
            // to the max Q seqlen if the caller did not set it.
            size_t kv_tokens = launch_params.total_kv_seqlen > 0
                ? static_cast<size_t>(launch_params.total_kv_seqlen)
                : params.b * params.s;
            size_t kv_size_in_bytes = kv_tokens * params.h_kv * params.d * 2 * {{ bytes_per_elt }};
            params.use_head_first_scheduling =
                (2 * kv_size_in_bytes <= launch_params.device_l2_cache_size);
        }

        block_size.x = 1;
        block_size.y = std::min(static_cast<int>(params.num_tiles), launch_params.multi_processor_count);
    } else {
        assert(false && "Invalid SCHEDULING_MODE");
    }

    // Reuse the same bytes_per_smem for launching kernels.
    constexpr int SMEM_BYTES = Ktraits::BYTES_PER_SMEM;
    if( launch_params.attention_mask_type == Attention_mask_type::PADDING ) {
{% if padding_mask %} // padding_mask
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ kernel_name }},
                                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                                         SMEM_BYTES));

        {{ kernel_name }}
            <<<block_size, Ktraits::THREADS, SMEM_BYTES, stream>>>({{ params_str }});
{% endif %} // padding_mask
    } else if( launch_params.attention_mask_type == Attention_mask_type::CAUSAL ) {
{% if causal_mask %} // causal_mask
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ causal_kernel_name }},
                                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                                         SMEM_BYTES));

        {{ causal_kernel_name }}
            <<<block_size, Ktraits::THREADS, SMEM_BYTES, stream>>>({{ params_str }});
{% endif %} // causal mask
    } else if( launch_params.attention_mask_type == Attention_mask_type::SLIDING_OR_CHUNKED_CAUSAL ) {
{% if sliding_or_chunked_causal_mask %} // sliding_or_chunked_causal_mask
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ sliding_or_chunked_causal_kernel_name }},
                                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                                         SMEM_BYTES));

        {{ sliding_or_chunked_causal_kernel_name }}
            <<<block_size, Ktraits::THREADS, SMEM_BYTES, stream>>>({{ params_str }});
{% endif %} // sliding_or_chunked_causal_mask
    } else if( launch_params.attention_mask_type == Attention_mask_type::CUSTOM_MASK ) {
{% if custom_mask %} // custom_mask
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ custom_mask_kernel_name }},
                                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                                         SMEM_BYTES));

        {{ custom_mask_kernel_name }}
            <<<block_size, Ktraits::THREADS, SMEM_BYTES, stream>>>({{ params_str }});
{% endif %} // custom mask
    }

}

#endif
{{ local_ns_close }}
