// 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); // // This one is slightly faster for causal? // // 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 A100 H100, 1st is faster with dropout, 3rd is faster without dropout // // For A6000, 1st is faster when causal, 3rd is faster when not causal // } template<> void run_mha_fwd_(Flash_fwd_params ¶ms, cudaStream_t stream) { run_mha_fwd_hdim192(params, stream); }