#include <flashinfer/attention/prefill.cuh>
#include "batch_prefill_config.inc"

namespace flashinfer {

constexpr auto use_custom_mask = {{ mask_mode }} == MaskMode::kCustom;

{# FA2DetermineCtaTileQ only selects {16, 32} for head_dim_vo >= 512 and
   {16, 64, 128} otherwise; don't instantiate unreachable variants. #}
{% for cta_tile_q in ([16, 32] if head_dim_vo | int >= 512 else [16, 64, 128]) %}
template cudaError_t BatchPrefillWithRaggedKVCacheDispatched<
    /*CTA_TILE_Q=*/{{cta_tile_q}}, {{head_dim_qk}}, {{head_dim_vo}}, {{pos_encoding_mode}}, {{use_fp16_qk_reduction}}, {{mask_mode}},
    {{ variant_name }}, RaggedParams>(RaggedParams params, {{ dtype_o }}* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream);
{% endfor %}

};  // namespace flashinfer
