dispatch_cutlassB Class — pytorch Architecture

Source Code

aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB.h lines 884–914

template <typename DT, typename T>
void dispatch_cutlassB(T cb, int cc = 0) {

    if (std::is_same_v<DT, cutlass::half_t> && 70 <= cc && cc < 75) {
        dispatch_cutlassB_f16_sm70(cb, cc);
    }
    if (std::is_same_v<DT, cutlass::bfloat16_t> && 80 <= cc && cc <= 120) {
        dispatch_cutlassB_bf16_sm80(cb, cc);
    }
    if (std::is_same_v<DT, cutlass::half_t> && 80 <= cc && cc <= 120) {
        dispatch_cutlassB_f16_sm80(cb, cc);
    }
    if (std::is_same_v<DT, cutlass::half_t> && 50 <= cc && cc < 70) {
        dispatch_cutlassB_f16_sm50(cb, cc);
    }
    if (std::is_same_v<DT, float> && 50 <= cc && cc < 70) {
        dispatch_cutlassB_f32_sm50(cb, cc);
    }
    if (std::is_same_v<DT, float> && 70 <= cc && cc < 75) {
        dispatch_cutlassB_f32_sm70(cb, cc);
    }
    if (std::is_same_v<DT, cutlass::half_t> && 75 <= cc && cc < 80) {
        dispatch_cutlassB_f16_sm75(cb, cc);
    }
    if (std::is_same_v<DT, float> && 75 <= cc && cc < 80) {
        dispatch_cutlassB_f32_sm75(cb, cc);
    }
    if (std::is_same_v<DT, float> && 80 <= cc && cc <= 120) {
        dispatch_cutlassB_f32_sm80(cb, cc);
    }
}