{{ copyright }}

//We can disable the FADD trick for archs with F2IP
{% if disable_fadd_trick %}
#ifdef USE_I2F_EMULATION_TRICK
#undef USE_I2F_EMULATION_TRICK
#endif

#ifdef USE_F2I_EMULATION_TRICK
#undef USE_F2I_EMULATION_TRICK
#endif
{% endif %}

#include <cuda.h>

#if CUDA_VERSION >= {{ min_cuda_version }}

#include <fused_multihead_attention_kernel_{{ kernel_variant }}.h>
{% if has_noloop %}
#include <fused_multihead_attention_kernel_{{ kernel_variant }}_noloop.h>
{% endif %}

{% if use_tma %}
// only included if tma is used.
#include <fmha/hopper/tma_descriptor.h>
{% endif %} //use_tma

{{ include_str }}
{{ local_ns_open }}
{{ bert_launch_params }}
{{ attn_mask_type_str }}

using Traits_p = fmha::{{ instruction_traits_p }};
using Traits_o = fmha::{{ instruction_traits_o }};

using Kernel_traits = {{ kernel_traits }}<
                       Traits_p,
                       Traits_o,
                       {{ seq_len }},
                       {{ head_size }},
                       {{ loop_step }},
                       {{ warps_m }},
                       {{ warps_n }},
                       2,
                       {{ kernel_flags }}>;

using Kernel_traits_causal = {{ kernel_traits }}<
                              Traits_p,
                              Traits_o,
                              {{ seq_len }},
                              {{ head_size }},
                              {{ loop_step }},
                              {{ warps_m }},
                              {{ warps_n }},
                              3,
                              {{ kernel_flags }}>;

using Kernel_traits_sliding_or_chunked_causal = {{ kernel_traits }}<
                                           Traits_p,
                                           Traits_o,
                                           {{ seq_len }},
                                           {{ head_size }},
                                           {{ loop_step }},
                                           {{ warps_m }},
                                           {{ warps_n }},
                                           4,
                                           {{ kernel_flags }}>;

{% if use_tma %} // use_tma

{% if padding_mask %} // padding_mask

extern "C"
__global__
void {{ kernel_name }}(const __grid_constant__ {{ params_type }} params){
  fused_multihead_attention::device_{{ kernel_variant }}_tma<Kernel_traits>(params);
}

{% endif %} // padding_mask

{% if causal_mask %} // causal_mask

extern "C"
__global__
void {{ causal_kernel_name }}(const __grid_constant__ {{ params_type }} params){
  fused_multihead_attention::device_{{ kernel_variant }}_tma<Kernel_traits_causal>(params);
}

{% endif %} // causal mask

{% if sliding_or_chunked_causal_mask %} // sliding_or_chunked_causal_mask

extern "C"
__global__
void {{ sliding_or_chunked_causal_kernel_name }}(const __grid_constant__ {{ params_type }} params){
  fused_multihead_attention::device_{{ kernel_variant }}_tma<Kernel_traits_sliding_or_chunked_causal>(params);
}

{% endif %} // sliding_or_chunked_causal_mask

{% else %}

{% if padding_mask %}

extern "C"
__global__
void {{ kernel_name }}(const __grid_constant__ {{ params_type }} params){
  fused_multihead_attention::device_{{ kernel_variant }}<Kernel_traits>(params);
}

{% endif %} // padding_mask

{% if causal_mask %} // causal_mask

extern "C"
__global__
void {{ causal_kernel_name }}(const __grid_constant__ {{ params_type }} params){
  fused_multihead_attention::device_{{ kernel_variant }}<Kernel_traits_causal>(params);
}

{% endif %} // causal mask

{% if sliding_or_chunked_causal_mask %} // sliding_or_chunked_causal_mask

extern "C"
__global__
void {{ sliding_or_chunked_causal_kernel_name }}(const __grid_constant__ {{ params_type }} params){
  fused_multihead_attention::device_{{ kernel_variant }}<Kernel_traits_sliding_or_chunked_causal>(params);
}
{% endif %}

{% endif %} // sliding_or_chunked_causal_mask

void {{ launcher_name }}({{ fused_multihead_attention_params_v2_str }} &params,
    const Launch_params &launch_params, cudaStream_t stream){
  // setting TMA descriptors if needed.
  // use_tma = {{ use_tma }}
{% if use_tma %}
    // declare TMA desc for Q, K, V
    typename fmha::Multiple_tma_descriptor<3> tma_desc_QKV;

    // GMEM pointers, the offset between each batch is d*3*h*seqlen
    // qkv pointer
    char *qkv_ptr = reinterpret_cast<char*>(params.qkv_ptr);

    // tensor size
    uint32_t tensor_size_qkv[3];
    tensor_size_qkv[2] = 1;
    tensor_size_qkv[1] = params.is_s_padded ? params.s * params.b : launch_params.seqlens[params.b];
    tensor_size_qkv[0] = (params.h + 2 * params.h_kv) * params.d;

    // box size for Q
    uint32_t box_size_q[3];
    box_size_q[2] = 1;
    box_size_q[1] = {{ loop_step }}; // STEP size
    box_size_q[0] = {{ head_size }}; // head_size

    // box size for k and v
    uint32_t box_size_kv[3];
    box_size_kv[2] = 1;
    box_size_kv[1] = params.s; // S, should not be actual_s, OOB will be filled with zeros.
    box_size_kv[0] = {{ head_size }}; // head_size

    // stride size
    uint64_t tensor_stride_qkv[2];
    tensor_stride_qkv[0] = tensor_size_qkv[0] * Traits_p::BITS_PER_ELEMENT_A / 8;
    tensor_stride_qkv[1] = tensor_size_qkv[1] * tensor_stride_qkv[0];

    // traversal stride
    uint32_t traversal_stride_qkv[3] = {1, 1, 1};

    // OOB fill zeros
    uint32_t oob_fill = 0;

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

    //setup the descriptors

    //setup the descriptor for Q
    tma_desc_QKV.set_tma_desctriptor(reinterpret_cast<void*>(qkv_ptr),
                                fmha::cudaTmaDescFormat::F16_RN, // tma format (data type). For now hardcode to fp16
                                fmha::cudaTmaDescInterleave::INTERLEAVE_DISABLED,
                                fmha::cudaTmaDescSwizzle::SWIZZLE_128B,
                                fmha::cudaTmaDescPromotion::PROMOTION_DISABLED,
                                tensor_size_qkv,
                                tensor_stride_qkv,
                                traversal_stride_qkv,
                                box_size_q,
                                oob_fill,
                                fp32_to_tf32,
                                &params.tma_desc_q);

    // setup the descriptor for K
    tma_desc_QKV.set_tma_desctriptor(reinterpret_cast<void*>(qkv_ptr),
                                fmha::cudaTmaDescFormat::F16_RN, // tma format (data type). For now hardcode to fp16
                                fmha::cudaTmaDescInterleave::INTERLEAVE_DISABLED,
                                fmha::cudaTmaDescSwizzle::SWIZZLE_128B,
                                fmha::cudaTmaDescPromotion::PROMOTION_DISABLED,
                                tensor_size_qkv,
                                tensor_stride_qkv,
                                traversal_stride_qkv,
                                box_size_kv,
                                oob_fill,
                                fp32_to_tf32,
                                &params.tma_desc_k);

    // setup the descriptor for V
    tma_desc_QKV.set_tma_desctriptor(reinterpret_cast<void*>(qkv_ptr),
                                fmha::cudaTmaDescFormat::F16_RN, // tma format (data type). For now hardcode to fp16
                                fmha::cudaTmaDescInterleave::INTERLEAVE_DISABLED,
                                fmha::cudaTmaDescSwizzle::SWIZZLE_128B,
                                fmha::cudaTmaDescPromotion::PROMOTION_DISABLED,
                                tensor_size_qkv,
                                tensor_stride_qkv,
                                traversal_stride_qkv,
                                box_size_kv,
                                oob_fill,
                                fp32_to_tf32,
                                &params.tma_desc_v);


{% endif %} // use_tma
  dim3 grid(params.h, params.b);
  // Use the same smem_size for all traits.
  constexpr int smem_size = Kernel_traits::BYTES_PER_SMEM;
  if( launch_params.attention_mask_type == Attention_mask_type::CAUSAL ) {
{% if causal_mask %} // causal_mask
    if( smem_size >= 48*1024 ) {
       FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ causal_kernel_name }},
                                        cudaFuncAttributeMaxDynamicSharedMemorySize,
                                        smem_size));
    }
    {{ causal_kernel_name }}<<<grid, Kernel_traits::THREADS, Kernel_traits::BYTES_PER_SMEM, 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
    if( smem_size >= 48*1024 ) {
       FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ sliding_or_chunked_causal_kernel_name }},
                                        cudaFuncAttributeMaxDynamicSharedMemorySize,
                                        smem_size));
    }
    {{ sliding_or_chunked_causal_kernel_name }}<<<grid, Kernel_traits::THREADS, Kernel_traits::BYTES_PER_SMEM, stream>>>({{ params_str }});
{% endif %} // sliding_or_chunked_causal_mask
  } else {
{% if padding_mask %} // padding_mask
    constexpr int smem_size = Kernel_traits::BYTES_PER_SMEM;
    if( smem_size >= 48*1024 ) {
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ kernel_name }},
                                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                                         smem_size));
    }
    {{ kernel_name }}<<<grid, Kernel_traits::THREADS, Kernel_traits::BYTES_PER_SMEM, stream>>>({{ params_str }});
{% endif %} // padding_mask
  }
}

{% if has_noloop %}


using Kernel_traits_nl = {{ kernel_traits }}<
                          Traits_p,
                          Traits_o,
                          {{ seq_len }},
                          {{ head_size }},
                          {{ noloop_step }},
                          {{ warps_m }},
                          {{ warps_n }},
                          2,
                          {{ kernel_flags }}>;

using Kernel_traits_causal_nl = {{ kernel_traits }}<
                                 Traits_p,
                                 Traits_o,
                                 {{ seq_len }},
                                 {{ head_size }},
                                 {{ noloop_step }},
                                 {{ warps_m }},
                                 {{ warps_n }},
                                 3,
                                 {{ kernel_flags }}>;

using Kernel_traits_sliding_or_chunked_causal_nl = {{ kernel_traits }}<
                                              Traits_p,
                                              Traits_o,
                                              {{ seq_len }},
                                              {{ head_size }},
                                              {{ noloop_step }},
                                              {{ warps_m }},
                                              {{ warps_n }},
                                              4,
                                              {{ kernel_flags }}>;

{% if padding_mask %} // padding_mask

extern "C"
__global__
void {{ kernel_name }}_nl({{ params_type }} params){
  fused_multihead_attention::device_{{ kernel_variant }}_nl<Kernel_traits_nl>(params);
}

{% endif %} // padding_mask

{% if causal_mask %} // causal_mask

extern "C"
__global__
void {{ causal_kernel_name }}_nl({{ params_type }} params){
  fused_multihead_attention::device_{{ kernel_variant }}_nl<Kernel_traits_causal_nl>(params);
}

{% endif %} // causal mask

{% if sliding_or_chunked_causal_mask %} // sliding_or_chunked_causal_mask

extern "C"
__global__
void {{ sliding_or_chunked_causal_kernel_name }}_nl({{ params_type }} params){
  fused_multihead_attention::device_{{ kernel_variant }}_nl<Kernel_traits_sliding_or_chunked_causal_nl>(params);
}

{% endif %} // sliding_or_chunked_causal_mask

void {{ launcher_name }}_nl({{ fused_multihead_attention_params_v2_str }} &params,
    const Launch_params& launch_params, cudaStream_t stream){
  constexpr int loop_iters = {{ seq_len }} / {{ noloop_step }};
  static_assert(loop_iters * {{ noloop_step }} == {{ seq_len }}, "");
  dim3 grid(params.h, params.b, loop_iters);

  // Use the same smem_size for all traits.
  constexpr int smem_size = Kernel_traits::BYTES_PER_SMEM;
  if( launch_params.attention_mask_type == Attention_mask_type::CAUSAL ) {
{% if causal_mask %} // causal_mask
    if( smem_size >= 48*1024 ) {
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ causal_kernel_name }}_nl,
                                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                                         smem_size));
    }
    {{ causal_kernel_name }}_nl<<<grid, Kernel_traits_nl::THREADS, Kernel_traits_nl::BYTES_PER_SMEM, 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
    if( smem_size >= 48*1024 ) {
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ sliding_or_chunked_causal_kernel_name }}_nl,
                                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                                         smem_size));
    }
    {{ sliding_or_chunked_causal_kernel_name }}_nl<<<grid, Kernel_traits_nl::THREADS, Kernel_traits_nl::BYTES_PER_SMEM, stream>>>({{ params_str }});
{% endif %} // sliding_or_chunked_causal_mask
  } else {
{% if padding_mask %} // padding_mask
    if( smem_size >= 48*1024 ) {
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ kernel_name }}_nl,
                                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                                         smem_size));
    }
    {{ kernel_name }}_nl<<<grid, Kernel_traits_nl::THREADS, Kernel_traits_nl::BYTES_PER_SMEM, stream>>>({{ params_str }});
{% endif %} // padding_mask
  }
}

{% endif %}

#else

void {{ launcher_name }}(const {{ params_type }} &params, cudaStream_t stream){
    assert(false && "Unsupported CUDA version");
}

{% if has_noloop %}

void {{ launcher_name }}_nl(const {{ params_type }} &params, cudaStream_t stream){
    assert(false && "Unsupported CUDA version");
}

{% endif %}

#endif
{{ local_ns_close }}
