File tree Expand file tree Collapse file tree
src/infiniop/ops/rms_norm/moore Expand file tree Collapse file tree Original file line number Diff line number Diff line change @@ -77,10 +77,14 @@ infiniStatus_t launchKernel(
7777
7878 if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F16 ) {
7979 LAUNCH_KERNEL (half , half , float );
80+ } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_BF16 ){
81+ LAUNCH_KERNEL (half , __mt_bfloat16 , float );
8082 } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F32 ) {
8183 LAUNCH_KERNEL (half , float , float );
8284 } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_BF16 ) {
8385 LAUNCH_KERNEL (__mt_bfloat16 , __mt_bfloat16 , float );
86+ } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F16 ) {
87+ LAUNCH_KERNEL (__mt_bfloat16 , half , float );
8488 } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F32 ) {
8589 LAUNCH_KERNEL (__mt_bfloat16 , float , float );
8690 } else if (atype == INFINI_DTYPE_F32 && wtype == INFINI_DTYPE_F32 ) {
You can’t perform that action at this time.
0 commit comments