// Copyright (c) 2023, Tri Dao. // Splitting the different head dimensions to different files to speed up compilation. #include "flash_fwd_launch_template.h" // template<> // void run_mha_fwd_(Flash_fwd_params ¶ms, cudaStream_t stream) { // using elem_type = cutlass::half_t; // BOOL_SWITCH(params.p_dropout < 1.f, Is_dropout, [&] { // run_flash_fwd, Is_dropout>(params, stream); // run_flash_fwd, Is_dropout>(params, stream); // run_flash_fwd, Is_dropout>(params, stream); // run_flash_fwd, Is_dropout>(params, stream); // // run_flash_fwd>(params, stream); // // run_flash_fwd>(params, stream); // // run_flash_fwd>(params, stream); // // run_flash_fwd>(params, stream); // // run_flash_fwd>(params, stream); // // For A6000, no-causal, 1st is fastest. causal, 4th is fastest. // // For A100, H100, 1st is fastest. // }); // } template<> void run_mha_fwd_(Flash_fwd_params ¶ms, cudaStream_t stream) { run_mha_fwd_hdim160(params, stream); }