{{ copyright }}

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

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

#include <cuda.h>
#include <stdexcept>

#if CUDA_VERSION >= {{ min_cuda_version }}


{% if not use_multi_cta %} // !use_multi_cta
{% if kernel_variant == "flash_attention" %}
#include <fused_multihead_flash_attention_kernel.h>
{% else %}
#include <fused_multihead_attention_kernel_{{ kernel_variant }}.h>
{% endif %}
{% endif %} // !use_multi_cta

{% if not use_multi_cta and has_noloop %} // !use_multi_cta && has_noloop
{% if kernel_variant == "flash_attention" %}
#include <fused_multihead_flash_attention_kernel_noloop.h>
{% if reload_q %} // reload_q: D > CTA_P_TILE_K, need tiled noloop kernel
#include <fused_multihead_flash_attention_kernel_noloop_tiled.h>
{% endif %} // reload_q
{% else %}
#include <fused_multihead_attention_kernel_1xN_noloop.h>
{% endif %}
{% endif %} // !use_multi_cta && has_noloop

{% if cross_mha %} // cross_mha
{% if has_noloop %} // has_noloop
#include <fused_multihead_cross_attention_kernel_1xN_noloop.h>
{% endif %}  // has_noloop
#include <fused_multihead_cross_attention_kernel_1xN.h>
{% endif %} // cross_mha

{% if use_multi_cta %} // use_multi_cta
#include <fused_multihead_attention_kernel_1xN_multi_cta.h>
{% endif %}

using Attention_mask_type = fmha::Attention_mask_type;
using Launch_params = bert::Fused_multihead_attention_launch_params;

{% if not cross_mha and not has_noloop %} // !cross_mha && !has_noloop (looped kernel)
using Kernel_traits = fmha::{{ kernel_traits }}<
    fmha::{{ instruction_traits }},
    {{ kv_loop_step }},
    {{ head_size }},
    {{ head_size_v }},
    {{ loop_step }},
    {{ warps_m }},
    {{ warps_n }},
    {{ ctas_per_head }},
    {{ kernel_flags }}>;

using Kernel_traits_causal = fmha::{{ kernel_traits }}<
    fmha::{{ instruction_traits }},
    {{ kv_loop_step }},
    {{ head_size }},
    {{ head_size_v }},
    {{ loop_step }},
    {{ warps_m }},
    {{ warps_n }},
    {{ ctas_per_head }},
    {{ kernel_flags }},
    /*causal mask*/ 3>;
{% endif %} // !cross_mha && !has_noloop

{% if not use_multi_cta and not cross_mha and not has_noloop %} // !use_multi_cta && !cross_mha && !has_noloop

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

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

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

  constexpr int smem_size = Kernel_traits::BYTES_PER_SMEM;
  if( launch_params.attention_mask_type == Attention_mask_type::CAUSAL ) {
    if( smem_size >= 48*1024 ) {
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ causal_kernel_name }},
                                            cudaFuncAttributeMaxDynamicSharedMemorySize,
                                            smem_size));
    }
    dim3 grid(params.h, params.b);
    {{ causal_kernel_name }}<<<grid, Kernel_traits_causal::THREADS, Kernel_traits_causal::BYTES_PER_SMEM, stream>>>(params);
  } else {
    if( smem_size >= 48*1024 ) {
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ kernel_name }},
                                            cudaFuncAttributeMaxDynamicSharedMemorySize,
                                            smem_size));
    }
    dim3 grid(params.h, params.b);
    {{ kernel_name }}<<<grid, Kernel_traits::THREADS, Kernel_traits::BYTES_PER_SMEM, stream>>>(params);
  }
}

{% endif %} // !use_multi_cta && !cross_mha && !has_noloop

{% if has_noloop and not cross_mha %} // has_noloop && !cross_mha (looped launcher stub — flash attention uses noloop only)

void {{ launcher_name }}(
    const {{ params_type }} &params,
    const Launch_params &launch_params,
    cudaStream_t stream){
    assert(false && "Flash attention uses noloop launchers");
}

{% endif %} // has_noloop && !cross_mha

{% if not use_multi_cta and has_noloop and not cross_mha %} // !use_multi_cta && has_noloop && !cross_mha

using Kernel_traits_nl = fmha::{{ kernel_traits }}<
    fmha::{{ instruction_traits }},
    {{ kv_loop_step }},
    {{ head_size }},
    {{ head_size_v }},
    {{ noloop_step }},
    1,
    {{ warps_m }} * {{ warps_n }},
    {{ ctas_per_head }},
    {{ kernel_flags }} | 0x200 /* no_loop flag */ >;

static_assert(Kernel_traits_nl::CTAS_PER_HEAD == 1, "");

using Kernel_traits_nl_causal = fmha::{{ kernel_traits }}<
    fmha::{{ instruction_traits }},
    {{ kv_loop_step }},
    {{ head_size }},
    {{ head_size_v }},
    {{ noloop_step }},
    1,
    {{ warps_m }} * {{ warps_n }},
    {{ ctas_per_head }},
    {{ kernel_flags }} | 0x200 /* no_loop flag */,
    /*causal mask*/ 3>;

static_assert(Kernel_traits_nl_causal::CTAS_PER_HEAD == 1, "");
static_assert(Kernel_traits_nl_causal::MASK_VERSION == 3, "");

{% if sliding_or_chunked_causal_mask == "1" %}
// MASK_VERSION 4 = sliding-or-chunked causal. Shared by the non-tiled (_nl) and
// non-reload tiled (_nl_tiled else-branch) launchers; without it both route
// SLIDING_OR_CHUNKED_CAUSAL to the maskless (dense) kernel.
using Kernel_traits_nl_sliding_or_chunked = fmha::{{ kernel_traits }}<
    fmha::{{ instruction_traits }},
    {{ kv_loop_step }},
    {{ head_size }},
    {{ head_size_v }},
    {{ noloop_step }},
    1,
    {{ warps_m }} * {{ warps_n }},
    {{ ctas_per_head }},
    {{ kernel_flags }} | 0x200 /* no_loop flag */,
    /*sliding-or-chunked-causal mask*/ 4>;

static_assert(Kernel_traits_nl_sliding_or_chunked::CTAS_PER_HEAD == 1, "");
static_assert(Kernel_traits_nl_sliding_or_chunked::MASK_VERSION == 4, "");
{% endif %}

extern "C"
__global__
void {{ kernel_name }}_nl({{ params_type }} params){
{% if kernel_variant == "flash_attention" %}
  fused_multihead_attention::device_flash_attention_nl<Kernel_traits_nl>(params);
{% else %}
  fused_multihead_attention::device_1xN_nl<Kernel_traits_nl>(params);
{% endif %}
}

extern "C"
__global__
void {{ causal_kernel_name }}_nl({{ params_type }} params){
{% if kernel_variant == "flash_attention" %}
  fused_multihead_attention::device_flash_attention_nl<Kernel_traits_nl_causal>(params);
{% else %}
  fused_multihead_attention::device_1xN_nl<Kernel_traits_nl_causal>(params);
{% endif %}
}

{% if sliding_or_chunked_causal_mask == "1" %}
extern "C"
__global__
void {{ sliding_or_chunked_causal_kernel_name }}_nl({{ params_type }} params){
{% if kernel_variant == "flash_attention" %}
  fused_multihead_attention::device_flash_attention_nl<Kernel_traits_nl_sliding_or_chunked>(params);
{% else %}
  fused_multihead_attention::device_1xN_nl<Kernel_traits_nl_sliding_or_chunked>(params);
{% endif %}
}
{% endif %}

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

  // Runtime loop_iters: flash attention supports variable sequence lengths.
  int loop_iters = ( params.s + {{ noloop_step }} - 1 )  / {{ noloop_step }};
  dim3 grid(loop_iters, params.h, params.b); // better locality
  if( launch_params.attention_mask_type == Attention_mask_type::CAUSAL ) {
    constexpr int smem_size = Kernel_traits_nl_causal::BYTES_PER_SMEM;
    if( smem_size >= 48*1024 ) {
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ causal_kernel_name }}_nl,
                                            cudaFuncAttributeMaxDynamicSharedMemorySize,
                                            smem_size));
    }
    {{ causal_kernel_name }}_nl<<<grid, Kernel_traits_nl_causal::THREADS, Kernel_traits_nl_causal::BYTES_PER_SMEM, stream>>>(params);
{% if sliding_or_chunked_causal_mask == "1" %}
  } else if( launch_params.attention_mask_type == Attention_mask_type::SLIDING_OR_CHUNKED_CAUSAL ) {
    constexpr int smem_size = Kernel_traits_nl_sliding_or_chunked::BYTES_PER_SMEM;
    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_sliding_or_chunked::THREADS, Kernel_traits_nl_sliding_or_chunked::BYTES_PER_SMEM, stream>>>(params);
{% endif %}
  } else {
    constexpr int smem_size = Kernel_traits_nl::BYTES_PER_SMEM;
    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);
  }
}

{% endif %} // !use_multi_cta && has_noloop && !cross_mha

{% if not use_multi_cta and tiled and not cross_mha %} // !use_multi_cta && tiled && !cross_mha
{% if reload_q %} // reload_q
// When D > CTA_P_TILE_K (RELOAD_Q=true), use the tiled noloop kernel which iterates
// over the K dimension in BMM1. This kernel requires ldgsts (cp.async) and
// WARPS_N=1 (no shared memory for softmax), so we define separate Kernel_traits
// with WARPS_M={{ warps_m }}, WARPS_N={{ warps_n }} instead of the nl layout (1, total).

using Kernel_traits_nl_tiled = fmha::{{ kernel_traits }}<
    fmha::{{ instruction_traits }},
    {{ kv_loop_step }},
    {{ head_size }},
    {{ head_size_v }},
    {{ noloop_step }},
    {{ warps_m }},
    {{ warps_n }},
    {{ ctas_per_head }},
    {{ kernel_flags }} | 0x200 /* no_loop flag */ >;

static_assert(Kernel_traits_nl_tiled::CTAS_PER_HEAD == 1, "");

using Kernel_traits_nl_tiled_causal = fmha::{{ kernel_traits }}<
    fmha::{{ instruction_traits }},
    {{ kv_loop_step }},
    {{ head_size }},
    {{ head_size_v }},
    {{ noloop_step }},
    {{ warps_m }},
    {{ warps_n }},
    {{ ctas_per_head }},
    {{ kernel_flags }} | 0x200 /* no_loop flag */,
    /*causal mask*/ 3>;

static_assert(Kernel_traits_nl_tiled_causal::CTAS_PER_HEAD == 1, "");
static_assert(Kernel_traits_nl_tiled_causal::MASK_VERSION == 3, "");

{% if sliding_or_chunked_causal_mask == "1" %}
// MASK_VERSION 4 = sliding-or-chunked causal (CAUSAL_MASK && SLIDING_WINDOW_ATTENTION).
// Without this variant the tiled-noloop launcher routes SLIDING_OR_CHUNKED_CAUSAL
// requests to the maskless (dense) kernel, silently producing bidirectional output.
using Kernel_traits_nl_tiled_sliding_or_chunked = fmha::{{ kernel_traits }}<
    fmha::{{ instruction_traits }},
    {{ kv_loop_step }},
    {{ head_size }},
    {{ head_size_v }},
    {{ noloop_step }},
    {{ warps_m }},
    {{ warps_n }},
    {{ ctas_per_head }},
    {{ kernel_flags }} | 0x200 /* no_loop flag */,
    /*sliding-or-chunked-causal mask*/ 4>;

static_assert(Kernel_traits_nl_tiled_sliding_or_chunked::CTAS_PER_HEAD == 1, "");
static_assert(Kernel_traits_nl_tiled_sliding_or_chunked::MASK_VERSION == 4, "");
{% endif %}

{% else %} // !reload_q
// Tiled (granular) variant: uses the same Kernel_traits_nl with USE_GRANULAR_TILING
// in the flags for correct tile decomposition (CTA_O_TILE_K computed from K_PER_MMA
// instead of S, which is 0 for flash attention). Routes to the regular noloop kernel
// which supports granular tiling without requiring ldgsts.
{% endif %}

extern "C"
__global__
void {{ kernel_name }}_nl_tiled({{ params_type }} params){
{% if kernel_variant == "flash_attention" %}
{% if reload_q %}
  fused_multihead_attention::device_flash_attention_nl_tiled<Kernel_traits_nl_tiled>(params);
{% else %}
  fused_multihead_attention::device_flash_attention_nl<Kernel_traits_nl>(params);
{% endif %}
{% else %}
  fused_multihead_attention::device_1xN_nl<Kernel_traits_nl>(params);
{% endif %}
}

extern "C"
__global__
void {{ causal_kernel_name }}_nl_tiled({{ params_type }} params){
{% if kernel_variant == "flash_attention" %}
{% if reload_q %}
  fused_multihead_attention::device_flash_attention_nl_tiled<Kernel_traits_nl_tiled_causal>(params);
{% else %}
  fused_multihead_attention::device_flash_attention_nl<Kernel_traits_nl_causal>(params);
{% endif %}
{% else %}
  fused_multihead_attention::device_1xN_nl<Kernel_traits_nl_causal>(params);
{% endif %}
}

{% if kernel_variant == "flash_attention" and sliding_or_chunked_causal_mask == "1" %}
extern "C"
__global__
void {{ sliding_or_chunked_causal_kernel_name }}_nl_tiled({{ params_type }} params){
{% if reload_q %}
  fused_multihead_attention::device_flash_attention_nl_tiled<Kernel_traits_nl_tiled_sliding_or_chunked>(params);
{% else %}
  fused_multihead_attention::device_flash_attention_nl<Kernel_traits_nl_sliding_or_chunked>(params);
{% endif %}
}
{% endif %}

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

  // Runtime loop_iters: flash attention supports variable sequence lengths.
  int loop_iters = ( params.s + {{ noloop_step }} - 1 )  / {{ noloop_step }};
  // When the V/output head dim exceeds the per-CTA output tile (CTA_O_TILE_N, capped at
  // 256 for granular tiling), the kernel splits the output feature dimension across
  // ctas_per_o_row CTAs (kernel: q_loop = blockIdx.x / ctas_per_o_row,
  // o_part = blockIdx.x % ctas_per_o_row). Each CTA recomputes BMM1+softmax and writes a
  // disjoint slice of O (no cross-CTA reduction), so grid.x must include this factor.
  // For head dims <= 256 this is 1, leaving existing kernels unchanged.
{% if reload_q %}
  constexpr int ctas_per_o_row =
      ( (int)Kernel_traits_nl_tiled::Cta_tile_o::VALID_N + (int)Kernel_traits_nl_tiled::Cta_tile_o::N - 1 )
      / (int)Kernel_traits_nl_tiled::Cta_tile_o::N;
{% else %}
  constexpr int ctas_per_o_row =
      ( (int)Kernel_traits_nl::Cta_tile_o::VALID_N + (int)Kernel_traits_nl::Cta_tile_o::N - 1 )
      / (int)Kernel_traits_nl::Cta_tile_o::N;
{% endif %}
  dim3 grid(loop_iters * ctas_per_o_row, params.h, params.b); // better locality
{% if reload_q %}
  if( launch_params.attention_mask_type == Attention_mask_type::CAUSAL ) {
    constexpr int smem_size = Kernel_traits_nl_tiled_causal::BYTES_PER_SMEM;
    if( smem_size >= 48*1024 ) {
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ causal_kernel_name }}_nl_tiled,
                                            cudaFuncAttributeMaxDynamicSharedMemorySize,
                                            smem_size));
    }
    {{ causal_kernel_name }}_nl_tiled<<<grid, Kernel_traits_nl_tiled_causal::THREADS, Kernel_traits_nl_tiled_causal::BYTES_PER_SMEM, stream>>>(params);
{% if sliding_or_chunked_causal_mask == "1" %}
  } else if( launch_params.attention_mask_type == Attention_mask_type::SLIDING_OR_CHUNKED_CAUSAL ) {
    constexpr int smem_size = Kernel_traits_nl_tiled_sliding_or_chunked::BYTES_PER_SMEM;
    if( smem_size >= 48*1024 ) {
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ sliding_or_chunked_causal_kernel_name }}_nl_tiled,
                                            cudaFuncAttributeMaxDynamicSharedMemorySize,
                                            smem_size));
    }
    {{ sliding_or_chunked_causal_kernel_name }}_nl_tiled<<<grid, Kernel_traits_nl_tiled_sliding_or_chunked::THREADS, Kernel_traits_nl_tiled_sliding_or_chunked::BYTES_PER_SMEM, stream>>>(params);
{% endif %}
  } else {
    constexpr int smem_size = Kernel_traits_nl_tiled::BYTES_PER_SMEM;
    if( smem_size >= 48*1024 ) {
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ kernel_name }}_nl_tiled,
                                            cudaFuncAttributeMaxDynamicSharedMemorySize,
                                            smem_size));
    }
    {{ kernel_name }}_nl_tiled<<<grid, Kernel_traits_nl_tiled::THREADS, Kernel_traits_nl_tiled::BYTES_PER_SMEM, stream>>>(params);
  }
{% else %}
  if( launch_params.attention_mask_type == Attention_mask_type::CAUSAL ) {
    constexpr int smem_size = Kernel_traits_nl_causal::BYTES_PER_SMEM;
    if( smem_size >= 48*1024 ) {
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ causal_kernel_name }}_nl_tiled,
                                            cudaFuncAttributeMaxDynamicSharedMemorySize,
                                            smem_size));
    }
    {{ causal_kernel_name }}_nl_tiled<<<grid, Kernel_traits_nl_causal::THREADS, Kernel_traits_nl_causal::BYTES_PER_SMEM, stream>>>(params);
{% if sliding_or_chunked_causal_mask == "1" %}
  } else if( launch_params.attention_mask_type == Attention_mask_type::SLIDING_OR_CHUNKED_CAUSAL ) {
    constexpr int smem_size = Kernel_traits_nl_sliding_or_chunked::BYTES_PER_SMEM;
    if( smem_size >= 48*1024 ) {
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ sliding_or_chunked_causal_kernel_name }}_nl_tiled,
                                            cudaFuncAttributeMaxDynamicSharedMemorySize,
                                            smem_size));
    }
    {{ sliding_or_chunked_causal_kernel_name }}_nl_tiled<<<grid, Kernel_traits_nl_sliding_or_chunked::THREADS, Kernel_traits_nl_sliding_or_chunked::BYTES_PER_SMEM, stream>>>(params);
{% endif %}
  } else {
    constexpr int smem_size = Kernel_traits_nl::BYTES_PER_SMEM;
    if( smem_size >= 48*1024 ) {
        FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ kernel_name }}_nl_tiled,
                                            cudaFuncAttributeMaxDynamicSharedMemorySize,
                                            smem_size));
    }
    {{ kernel_name }}_nl_tiled<<<grid, Kernel_traits_nl::THREADS, Kernel_traits_nl::BYTES_PER_SMEM, stream>>>(params);
  }
{% endif %}
}

{% endif %} // !use_multi_cta && tiled && !cross_mha

{% if cross_mha %} // cross_mha
{% if not use_multi_cta and has_noloop %} // !use_multi_cta && has_noloop

using Kernel_traits_nl = fmha::{{ kernel_traits }}<
    fmha::{{ instruction_traits }},
    {{ seq_len }},
    {{ head_size }},
    {{ head_size_v }},
    {{ noloop_step }},
    1,
    {{ warps_m }} * {{ warps_n }},
    {{ ctas_per_head }},
    {{ kernel_flags }}>;

static_assert(Kernel_traits_nl::CTAS_PER_HEAD == 1, "");

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

void {{ launcher_name }}_nl(
    const {{ params_type }} &params,
    // const Launch_params &launch_params, // TODO
    cudaStream_t stream){

  constexpr int smem_size = Kernel_traits_nl::BYTES_PER_SMEM;
  if( smem_size >= 48*1024 ) {
    FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ kernel_name }}_nl,
                                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                                         smem_size));
  }
  const int loop_iters = (params.s_q + {{ noloop_step }}-1) / {{ noloop_step }};
  // if (loop_iters * {{ noloop_step }} != params.s_q) {
  //   throw std::runtime_error("Incorrect seq len -- loop_iters * noloop_step != params.s_q");
  // }
  assert(loop_iters * {{ noloop_step }} >= params.s_q);
  dim3 grid(params.h, params.b, loop_iters);
  {{ kernel_name }}_nl<<<grid, Kernel_traits_nl::THREADS, Kernel_traits_nl::BYTES_PER_SMEM, stream>>>(params);
}

{% endif %} // !use_multi_cta && has_noloop

{% if not use_multi_cta %} // !use_multi_cta

using Kernel_traits = fmha::{{ kernel_traits }}<
    fmha::{{ instruction_traits }},
    {{ seq_len }},
    {{ head_size }},
    {{ head_size_v }},
    {{ loop_step }},
    {{ warps_m }},
    {{ warps_n }},
    {{ ctas_per_head }},
    {{ kernel_flags }}>;

extern "C"
__global__
void {{ kernel_name }}({{ params_type }} params){
  fused_multihead_attention::device_mhca_1xN<Kernel_traits>(params);
}

void {{ launcher_name }}(
    const {{ params_type }} &params,
    // const Launch_params &launch_params, // TODO
    cudaStream_t stream){

  constexpr int smem_size = Kernel_traits::BYTES_PER_SMEM;
  if( smem_size >= 48*1024 ) {
    FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ kernel_name }},
                                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                                         smem_size));
  }
  dim3 grid(params.h, params.b);
  {{ kernel_name }}<<<grid, Kernel_traits::THREADS, Kernel_traits::BYTES_PER_SMEM, stream>>>(params);
}

{% endif %} // !use_multi_cta

{% endif %} // cross_mha

{% if use_multi_cta %} // use_multi_cta

// If that assert gets triggered - increase the value of MAX_STGS_PER_LOOP in "setup.py".
static_assert(Kernel_traits::Gmem_tile_o::STGS_PER_LOOP <= {{ MAX_STGS_PER_LOOP }}, "");

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

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

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

  assert(params.heads_per_wave != 0 && "Heads per wave is not set, but multi cta is requested");

  // Clear the barriers and locks.
  cudaMemsetAsync(params.counters, 0, 3*params.heads_per_wave*sizeof(int), stream);

  // We may use more than 48kB of shared memory.
  if( launch_params.attention_mask_type == Attention_mask_type::CAUSAL ) {
    constexpr int smem_size = Kernel_traits_causal::BYTES_PER_SMEM;
    if( smem_size >= 48*1024 ) {
      FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ causal_kernel_name }},
                                          cudaFuncAttributeMaxDynamicSharedMemorySize,
                                          smem_size));
    }
    // Launch one wave.
    dim3 grid(Kernel_traits_causal::CTAS_PER_HEAD, params.heads_per_wave), block(Kernel_traits_causal::THREADS);
    void *params_ = (void*) &params;
    FMHA_CHECK_CUDA(cudaLaunchCooperativeKernel((void*) &{{ causal_kernel_name }}, grid, block, (void**) &params_, smem_size, stream));
  } else {
    constexpr size_t smem_size = Kernel_traits::BYTES_PER_SMEM;
    if( smem_size >= 48*1024 ) {
      FMHA_CHECK_CUDA(cudaFuncSetAttribute({{ kernel_name }},
                                          cudaFuncAttributeMaxDynamicSharedMemorySize,
                                          smem_size));
    }
    // Launch one wave.
    dim3 grid(Kernel_traits::CTAS_PER_HEAD, params.heads_per_wave), block(Kernel_traits::THREADS);
    void *params_ = (void*) &params;
    FMHA_CHECK_CUDA(cudaLaunchCooperativeKernel((void*) &{{ kernel_name }}, grid, block, (void**) &params_, smem_size, stream));
  }
}

{% endif %} // use_multi_cta

void {{ launcher_name }}_get_max_heads_per_wave(int *heads_per_wave) {
{% if use_multi_cta %} // use_multi_cta
    // Determine the number of SMs and CTAs.
    int dev;
    cudaGetDevice(&dev);
    cudaDeviceProp props;
    FMHA_CHECK_CUDA(cudaGetDeviceProperties(&props, dev));

    // The number of CTAs per SM.
    constexpr size_t smem_size = Kernel_traits::BYTES_PER_SMEM;
    int ctas_per_sm;
    FMHA_CHECK_CUDA(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&ctas_per_sm,
                                                            &{{ kernel_name }},
                                                            Kernel_traits::THREADS,
                                                            smem_size));

    // The number of heads per wave.
    *heads_per_wave = props.multiProcessorCount * ctas_per_sm / Kernel_traits::CTAS_PER_HEAD;
{% else %} // use_multi_cta
    *heads_per_wave = 0;
{% endif %} // use_multi_cta
}

#else // CUDA_VERSION >= {{ min_cuda_version }}

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

{% if has_noloop %} // has_noloop

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

{% endif %} // has_noloop

#endif // CUDA_VERSION >= {{ min_cuda_version }}
