fmha_bwd_hdim32.cu 723 Bytes
Newer Older
1
2
// Copyright (c) 2022, Tri Dao.

Tri Dao's avatar
Tri Dao committed
3
// Splitting the different head dimensions to different files to speed up compilation.
4
5
6
7

#include "fmha_bwd_launch_template.h"

void run_fmha_bwd_hdim32(FMHA_dgrad_params &params, cudaStream_t stream, const bool configure) {
8
    FP16_SWITCH(params.is_bf16, ({
9
10
11
12
13
14
15
        if (params.seqlen_k == 128) {
            using Kernel_traits = FMHA_kernel_traits<128, 32, 16, 1, 8, 0x08u, elem_type>;
            run_fmha_bwd_loop<Kernel_traits>(params, stream, configure);
        } else if (params.seqlen_k >= 256) {
            using Kernel_traits = FMHA_kernel_traits<256, 32, 16, 1, 8, 0x08u, elem_type>;
            run_fmha_bwd_loop<Kernel_traits>(params, stream, configure);
        }
16
    }));
17
}