pod_kernel_inst.jinja

#include <flashinfer/attention/pod.cuh>
#include "pod_config.inc"

namespace flashinfer {

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

{% for cta_tile_q_p in [16, 64, 128] %}
{% for cta_tile_q_d in [16, 64, 128] %}
template cudaError_t PODWithPagedKVCacheDispatched<
    /*CTA_TILE_Q_P=*/{{cta_tile_q_p}}, /*CTA_TILE_Q_D=*/{{cta_tile_q_d}}, {{head_dim_qk}}, {{head_dim_vo}}, {{pos_encoding_mode}}, {{use_fp16_qk_reduction}}, {{mask_mode}},
    {{ variant_name }}, PagedParams, PagedParams>(PagedParams prefill_params, PagedParams decode_params, {{ dtype_o }}* tmp_v, float* tmp_s, cudaStream_t stream);
{% endfor %}
{% endfor %}

};  // namespace flashinfer