diff --git a/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm2_blockscale.hpp b/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm2_blockscale.hpp index c6246db583..e68fbba50c 100644 --- a/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm2_blockscale.hpp +++ b/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm2_blockscale.hpp @@ -6,6 +6,7 @@ #include #include #include +#include #include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp" #include "ck/tensor_operation/gpu/device/device_base.hpp" @@ -77,6 +78,8 @@ struct ReferenceMoeGemm2BlockScale : public device::BaseOperator float Run(const Argument& arg) { + std::vector n_locks(arg.c_t_n_.mDesc.GetLengths()[1]); + arg.c_t_n_.SetZero(); auto f_mk_kn_mn = [&](auto m, auto n) { const int K = arg.a_t_k_k_.mDesc.GetLengths()[2]; @@ -142,6 +145,7 @@ struct ReferenceMoeGemm2BlockScale : public device::BaseOperator { arg.c_element_op_(v_c, v_acc, 1.f); } + std::lock_guard lock(n_locks[n]); arg.c_t_n_(t, n) += v_c; } };