From 4f4bce30cffac7962e43223651d367a138226ceb Mon Sep 17 00:00:00 2001 From: root Date: Thu, 27 Mar 2025 08:56:20 +0000 Subject: [PATCH] fuse gelu silu act in moe gemm1 --- .../moe_gemm1_xdl_fp8.cpp | 29 ++- .../gpu/device/impl/device_moe_gemm.hpp | 12 +- .../gpu/grid/gridwise_moe_gemm.hpp | 222 +++++++++++------- .../cpu/reference_moe_gemm.hpp | 62 +++-- 4 files changed, 199 insertions(+), 126 deletions(-) diff --git a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8.cpp b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8.cpp index e7b6816afd..6c229a3e8e 100644 --- a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8.cpp +++ b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8.cpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. #include #include @@ -81,12 +81,10 @@ struct MulABScale } }; - // using DsLayout = DsLayoutGate; // using DsDataType = DsDataTypeGate; using CDEElementOp = MulABScale; - void preShuffleBuffer(const B0DataType* src, B0DataType* dst, int N, int K, int NXdl) { int KPack = 16 / sizeof(B0DataType); @@ -136,6 +134,7 @@ static constexpr ck::index_t BK1 = 16 / sizeof(B0DataType); static constexpr ck::index_t EVec = 16 / sizeof(EDataType); static constexpr ck::index_t D0Vec = 1; static constexpr ck::index_t D1Vec = 1; +static constexpr ck::index_t ActOP = 2; // using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultiD_Xdl_CShuffle_V3 using DeviceOpInstance = ck::tensor_operation::device::DeviceMoeGemm // clang-format off @@ -156,7 +155,7 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceMoeGemm // MXdlPerWave| NXdlPerWave| _MBlock_MWaveMPerXdl| ScalarPerVector| // PerShuffle| PerShuffle| _NBlock_NWaveNPerXdl| _NWaveNPerXdl| 1, 1, S<1, 32, 1, 8>, S, - ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, Nswizzle, true, true, int32_t, A0DataType>; + ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, ActOP, Nswizzle, true, true, int32_t, A0DataType>; // clang-format on @@ -263,8 +262,10 @@ int main(int argc, char* argv[]) Tensor b0_e_n_k(HostTensorDescriptor({experts, K, N * 2}, {N * 2 * K, 1, K})); Tensor b0_preshuffled(HostTensorDescriptor({experts, K, N * 2}, {N * 2 * K, 1, K})); Tensor d0_t_n(HostTensorDescriptor({tokens, N}, {StrideDs[0], 0})); - // Tensor d1_e_n(HostTensorDescriptor({experts, N * 2}, {StrideDs[1] ? StrideDs[1] * N * 2: 1, StrideDs[1]})); - Tensor d1_e_n(HostTensorDescriptor({experts, N * 2}, {StrideDs[1] * N * 2, StrideDs[1]})); + // Tensor d1_e_n(HostTensorDescriptor({experts, N * 2}, {StrideDs[1] ? StrideDs[1] * + // N * 2: 1, StrideDs[1]})); + Tensor d1_e_n( + HostTensorDescriptor({experts, N * 2}, {StrideDs[1] * N * 2, StrideDs[1]})); Tensor e_t_n_host_result(HostTensorDescriptor({tokens, topk, N}, {topk * N, N, 1})); Tensor e_t_n_device_result( HostTensorDescriptor({tokens, topk, N}, {topk * N, N, 1})); @@ -278,10 +279,10 @@ int main(int argc, char* argv[]) { case 0: break; case 1: - a0_t_k.GenerateTensorValue(GeneratorTensor_2{0, 2}); - b0_e_n_k.GenerateTensorValue(GeneratorTensor_2{0, 2}); - d0_t_n.GenerateTensorValue(GeneratorTensor_2{0, 2}); - d1_e_n.GenerateTensorValue(GeneratorTensor_2{0, 2}); + a0_t_k.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); + b0_e_n_k.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); + d0_t_n.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); + d1_e_n.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); break; case 2: a0_t_k.GenerateTensorValue(GeneratorTensor_1{}); @@ -329,7 +330,8 @@ int main(int argc, char* argv[]) int NPerXdl = device_op.GetPreShuffleParameters(); - preShuffleBuffer(b0_e_n_k.mData.data(), b0_preshuffled.mData.data(), N * 2 * experts, K, NPerXdl); + preShuffleBuffer( + b0_e_n_k.mData.data(), b0_preshuffled.mData.data(), N * 2 * experts, K, NPerXdl); b0_device_buf.ToDevice(b0_preshuffled.mData.data()); @@ -394,7 +396,8 @@ int main(int argc, char* argv[]) AccDataType, PassThrough, PassThrough, - PassThrough>; + PassThrough, + ActOP>; auto ref_moe_gemm = ReferenceGemmInstance{}; auto ref_invoker = ref_moe_gemm.MakeInvoker(); @@ -437,7 +440,7 @@ int main(int argc, char* argv[]) e_t_n_device_result.savetxt("out.txt"); e_t_n_host_result.savetxt("ref.txt"); return ck::utils::check_err( - e_t_n_device_result, e_t_n_host_result, "Error: Incorrect results!", 1e-3, 5e-2) + e_t_n_device_result, e_t_n_host_result, "Error: Incorrect results!", 1e-3, 5e-1) ? 0 : 1; } diff --git a/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm.hpp b/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm.hpp index efe5c53d2d..096011051e 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm.hpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. #pragma once @@ -65,6 +65,7 @@ template ( - karg.p_sorted_token_ids, - karg.p_sorted_expert_ids, - karg.p_max_token_id, - karg.p_a_grid + splitk_batch_offset.a_k_split_offset, - karg.p_b_grid + splitk_batch_offset.b_k_split_offset, - karg.p_ds_grid, - karg.p_c_grid, - p_shared, - p_shared1, - karg, - karg.a_element_op, - karg.b_element_op, - karg.c_element_op); + GridwiseGemm::template Run_2Lds( + karg.p_sorted_token_ids, + karg.p_sorted_expert_ids, + karg.p_max_token_id, + karg.p_a_grid + splitk_batch_offset.a_k_split_offset, + karg.p_b_grid + splitk_batch_offset.b_k_split_offset, + karg.p_ds_grid, + karg.p_c_grid, + p_shared, + p_shared1, + karg, + karg.a_element_op, + karg.b_element_op, + karg.c_element_op); #else ignore = karg; #endif // end of if (defined(__gfx9__)) @@ -145,6 +152,7 @@ template - __host__ __device__ static auto - MakeCGridDescriptor_M_N(IndexType M, IndexType MPad, IndexType N, IndexType NPad, IndexType StrideC) + __host__ __device__ static auto MakeCGridDescriptor_M_N( + IndexType M, IndexType MPad, IndexType N, IndexType NPad, IndexType StrideC) { const auto c_grid_desc_mraw_nraw = [&]() { if constexpr(is_same::value) @@ -1222,7 +1230,8 @@ struct GridwiseMoeGemm } gather_offsets(m0) = static_cast(token_offset) * problem.K; }); - const index_t expert_stride = __builtin_amdgcn_readfirstlane(problem.N * problem.K * (IsInputGemm ? 2 : 1)); + const index_t expert_stride = + __builtin_amdgcn_readfirstlane(problem.N * problem.K * (IsInputGemm ? 2 : 1)); // N0, K0, Blocksize*KPack const index_t n_block_data_idx_on_grid = @@ -1309,13 +1318,13 @@ struct GridwiseMoeGemm const index_t num_k_block_main_loop = __builtin_amdgcn_readfirstlane( (a_grid_desc_ak0_m_ak1.GetLength(I0) * a_grid_desc_ak0_m_ak1.GetLength(I2)) / KPerBlock); - if constexpr (IsInputGemm) + if constexpr(IsInputGemm) { const BDataType* p_b_grid_up = p_b_grid + expert_stride / 2; - const auto b_grid_buf_up = make_dynamic_buffer( + const auto b_grid_buf_up = make_dynamic_buffer( p_b_grid_up + expert_id * expert_stride / BPackedSize, b_grid_desc_bpreshuffled.GetElementSpaceSize()); - auto b_blockwise_copy_up = ThreadwiseTensorSliceTransfer_v2< + auto b_blockwise_copy_up = ThreadwiseTensorSliceTransfer_v2< BDataType, BDataType, decltype(b_grid_desc_bpreshuffled), @@ -1326,44 +1335,46 @@ struct GridwiseMoeGemm BBlockTransferSrcScalarPerVector, BThreadTransferSrcResetCoordinateAfterRun, true>(b_grid_desc_bpreshuffled, - make_multi_index(n_block_data_idx_on_grid, - get_warp_local_1d_id() % NWave, - 0, - KPack * (get_thread_local_1d_id() % warpSize))); - blockwise_gemm_pipeline.template Run(a_grid_desc_ak0_m_ak1, - a_block_desc_ak0_m_ak1, - a_blockwise_copy, - a_grid_buf, - a_block_buf, - a_block_slice_copy_step, - b_grid_desc_bpreshuffled, - b_blockwise_copy, - b_blockwise_copy_up, - b_grid_buf, - b_grid_buf_up, - b_block_buf, - b_block_slice_copy_step, - c_thread_buf, - c_thread_buf_up, - num_k_block_main_loop); + make_multi_index(n_block_data_idx_on_grid, + get_warp_local_1d_id() % NWave, + 0, + KPack * (get_thread_local_1d_id() % warpSize))); + blockwise_gemm_pipeline.template Run( + a_grid_desc_ak0_m_ak1, + a_block_desc_ak0_m_ak1, + a_blockwise_copy, + a_grid_buf, + a_block_buf, + a_block_slice_copy_step, + b_grid_desc_bpreshuffled, + b_blockwise_copy, + b_blockwise_copy_up, + b_grid_buf, + b_grid_buf_up, + b_block_buf, + b_block_slice_copy_step, + c_thread_buf, + c_thread_buf_up, + num_k_block_main_loop); } else { - blockwise_gemm_pipeline.template Run(a_grid_desc_ak0_m_ak1, - a_block_desc_ak0_m_ak1, - a_blockwise_copy, - a_grid_buf, - a_block_buf, - a_block_slice_copy_step, - b_grid_desc_bpreshuffled, - b_blockwise_copy, - b_grid_buf, - b_block_buf, - b_block_slice_copy_step, - c_thread_buf, - num_k_block_main_loop); + blockwise_gemm_pipeline.template Run( + a_grid_desc_ak0_m_ak1, + a_block_desc_ak0_m_ak1, + a_blockwise_copy, + a_grid_buf, + a_block_buf, + a_block_slice_copy_step, + b_grid_desc_bpreshuffled, + b_blockwise_copy, + b_grid_buf, + b_block_buf, + b_block_slice_copy_step, + c_thread_buf, + num_k_block_main_loop); } - + // shuffle C and write out { static_assert(MXdlPerWave % CShuffleMXdlPerWavePerShuffle == 0 && @@ -1371,7 +1382,7 @@ struct GridwiseMoeGemm "wrong!"); constexpr index_t MWave = MPerBlock / (MXdlPerWave * MPerXdl); - + // TODO: hacky, fix it! constexpr auto c_thread_desc_m0_n0_m1_n1_m2_m3_m4_n2 = blockwise_gemm_pipeline.GetCThreadDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(); @@ -1389,14 +1400,17 @@ struct GridwiseMoeGemm constexpr auto M3 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I5); constexpr auto M4 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I6); constexpr auto N2 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I7); - + // mul scales - const float *p_scale_b = p_ds_grid[I1]; - if constexpr (PerTokenQuant) + const float* p_scale_b = p_ds_grid[I1]; + if constexpr(PerTokenQuant) { constexpr index_t scale_stride = (IsInputGemm ? 2 : 1); - p_scale_b += expert_id * problem.N * scale_stride + block_n_id * NPerBlock + get_warp_local_1d_id() % NWave * NPerXdl + threadIdx.x % NPerXdl; - } else { + p_scale_b += expert_id * problem.N * scale_stride + block_n_id * NPerBlock + + get_warp_local_1d_id() % NWave * NPerXdl + threadIdx.x % NPerXdl; + } + else + { p_scale_b += expert_id; } const float* p_sorted_weights_0 = p_ds_grid[I0]; @@ -1404,27 +1418,34 @@ struct GridwiseMoeGemm static_assert(M4 == 4); const index_t m1 = get_warp_local_1d_id() / NWave; const index_t m3 = threadIdx.x % get_warp_size() / MPerXdl; - vector_type scale_token_ids; - vector_type topk_weights; // for gemm2 only + vector_type scale_token_ids; + vector_type topk_weights; // for gemm2 only static_for<0, NXdlPerWave, 1>{}([&](auto n0) { const float scale_b = p_scale_b[n0 * NWave * PerTokenQuant]; - static_for<0, MXdlPerWave, 1>{}([&](auto m0) { // MXDLPerWave - static_for<0, M2, 1>{}([&](auto m2) { // m_inst_num_groups_per_blk - const index_t m_pos = block_m_id * MPerBlock + m0 * M1 * M2 * M3 * M4 + m1 * M2 * M3 * M4 + m2 * M3 * M4 + m3 * M4; - if constexpr(PerTokenQuant) { - scale_token_ids = *c_style_pointer_cast *>(p_sorted_token_ids + m_pos); - } - if constexpr (!IsInputGemm) + static_for<0, MXdlPerWave, 1>{}([&](auto m0) { // MXDLPerWave + static_for<0, M2, 1>{}([&](auto m2) { // m_inst_num_groups_per_blk + const index_t m_pos = block_m_id * MPerBlock + m0 * M1 * M2 * M3 * M4 + + m1 * M2 * M3 * M4 + m2 * M3 * M4 + m3 * M4; + if constexpr(PerTokenQuant) { - topk_weights = *c_style_pointer_cast *>(p_ds_grid[I2] + m_pos); + scale_token_ids = + *c_style_pointer_cast*>( + p_sorted_token_ids + m_pos); } - static_for<0, M4, 1>{}([&](auto m4) { // m_inst_group_size + if constexpr(!IsInputGemm) + { + topk_weights = *c_style_pointer_cast*>( + p_ds_grid[I2] + m_pos); + } + static_for<0, M4, 1>{}([&](auto m4) { // m_inst_group_size float scale_a = [&]() { if constexpr(PerTokenQuant) { index_t fused_token = scale_token_ids.AsType()[m4]; - const index_t token_offset = fused_token & 0xffffff; - return token_offset < problem.NumTokens ? p_sorted_weights_0[token_offset] : 0.0; + const index_t token_offset = fused_token & 0xffffff; + return token_offset < problem.NumTokens + ? p_sorted_weights_0[token_offset] + : 0.0; } else { @@ -1432,19 +1453,36 @@ struct GridwiseMoeGemm } }(); constexpr index_t c_offset = - blockwise_gemm_pipeline.GetCThreadDesc().CalculateOffset(make_tuple(m0, n0, m2 * M4 + m4)); + blockwise_gemm_pipeline.GetCThreadDesc().CalculateOffset( + make_tuple(m0, n0, m2 * M4 + m4)); constexpr auto cidx = Number{}; - if constexpr (IsInputGemm) // gu fusion + if constexpr(IsInputGemm) // gu fusion { - const float scale_up = p_scale_b[(n0 * NPerXdl + problem.N) * PerTokenQuant]; - auto gate = scale_a * scale_b * c_thread_buf[cidx]; - auto up = scale_a * scale_up * c_thread_buf_up[cidx]; - gate = gate * math::rcp(1.0 + math::exp(-gate)); - c_thread_buf(cidx) = gate * up; - } - else + if constexpr(ActivationOperation == Activation::silu) + { + tensor_operation::element_wise::Silu{}(c_thread_buf(cidx), + c_thread_buf(cidx)); + } + else if(ActivationOperation == Activation::gelu) + { + tensor_operation::element_wise::Gelu{}(c_thread_buf(cidx), + c_thread_buf(cidx)); + } + else if(ActivationOperation == Activation::swiglu) + { + const float scale_up = + p_scale_b[(n0 * NPerXdl + problem.N) * PerTokenQuant]; + auto gate = scale_a * scale_b * c_thread_buf[cidx]; + auto up = scale_a * scale_up * c_thread_buf_up[cidx]; + gate = gate * math::rcp(1.0 + math::exp(-gate)); + c_thread_buf(cidx) = gate * up; + } + } + else { - c_thread_buf(cidx) = scale_a * scale_b * topk_weights.AsType()[m4] * c_thread_buf[cidx]; + c_thread_buf(cidx) = scale_a * scale_b * + topk_weights.AsType()[m4] * + c_thread_buf[cidx]; } }); }); @@ -1612,7 +1650,7 @@ struct GridwiseMoeGemm Sequence, uniform_sequence_gen_t>, // ThreadTransferSrcResetCoordinateAfterRunFlags - Sequence, // ThreadTransferDstResetCoordinateAfterRunFlags + Sequence, // ThreadTransferDstResetCoordinateAfterRunFlags IndexType, 1, // ScatterDim true, // OutputScatter: false, only use scatter weights @@ -1623,7 +1661,8 @@ struct GridwiseMoeGemm make_tuple(make_multi_index(0, 0, block_n_id, 0)), c_element_op}; - auto c_grid_buf = make_dynamic_buffer(p_c_grid, c_grid_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize()); + auto c_grid_buf = make_dynamic_buffer( + p_c_grid, c_grid_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize()); constexpr auto sfc_c_vgpr = SpaceFillingCurve, Sequence<0, 1, 2, 3, 4, 5, 6, 7>, @@ -1662,7 +1701,7 @@ struct GridwiseMoeGemm block_m_id * MPerBlock + threadIdx.x / ENThreads * EMRepeats + dstidx(I1); static_for<0, EMRepeats, 1>{}([&](auto m0) { const index_t fused_token = p_sorted_token_ids[c_token_pos + m0]; - IndexType token_offset = fused_token & 0xffffff; + IndexType token_offset = fused_token & 0xffffff; if constexpr(IsInputGemm) { token_offset = token_offset * problem.TopK + (fused_token >> 24); @@ -2039,7 +2078,8 @@ struct GridwiseMoeGemm // if(i.value == 1) // { // ptr_ += - // expert_id * (problem.StrideDs[1] ? problem.StrideDs[1] * problem.N : 1); + // expert_id * (problem.StrideDs[1] ? problem.StrideDs[1] * problem.N : + // 1); // } return make_dynamic_buffer( ptr_, ds_grid_desc_m_n[i].GetElementSpaceSize()); @@ -2105,7 +2145,7 @@ struct GridwiseMoeGemm Sequence, uniform_sequence_gen_t>, // ThreadTransferSrcResetCoordinateAfterRunFlags - Sequence, // ThreadTransferDstResetCoordinateAfterRunFlags + Sequence, // ThreadTransferDstResetCoordinateAfterRunFlags IndexType, 1, // ScatterDim true, // OutputScatter: false, only use scatter weights @@ -2150,7 +2190,7 @@ struct GridwiseMoeGemm CDEBlockTransferCluster{}.At(I2) * CDEBlockTransferCluster{}.At(I3); static_for<0, num_access, 1>{}([&](auto access_id) { // make sure it's safe to write to LDS - StaticallyIndexedArray scatter_offsets; + StaticallyIndexedArray scatter_offsets; auto dstidx = sfc_cde_block.GetIndex(access_id); const index_t c_token_pos = diff --git a/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm.hpp b/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm.hpp index 2fba458c52..8c50923912 100644 --- a/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm.hpp +++ b/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm.hpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. #pragma once @@ -22,11 +22,13 @@ template + index_t ActivationType_ = 0, + typename ComputeTypeA = CDataType, + typename ComputeTypeB = ComputeTypeA> struct ReferenceMoeGemm : public device::BaseOperator { // Argument + static constexpr auto ActivationType = ActivationType_; struct Argument : public device::BaseArgument { Argument(const Tensor& sorted_token_ids, @@ -78,14 +80,20 @@ struct ReferenceMoeGemm : public device::BaseOperator float Run(const Argument& arg) { + if constexpr(ActivationType > 2) + { + static_assert(false, "Not supported activation type"); + } const int full_n = arg.c_t_k_n_.mDesc.GetLengths()[2]; - auto f_mk_kn_mn = [&](auto m, auto n) { + auto f_mk_kn_mn = [&](auto m, auto n) { const int K = arg.a_t_k_.mDesc.GetLengths()[1]; - AccDataType v_acc{0}; AccDataType v_acc_up{0}; + ComputeTypeB v_b_up{0}; + AccDataType v_acc{0}; + ComputeTypeA v_a{0}; ComputeTypeB v_b{0}; - ComputeTypeB v_b_up{0}; + const int t = arg.sorted_token_ids_(m) & 0xffffff; const int topk_id = (arg.sorted_token_ids_(m) & 0xff000000) >> 24; const int e = arg.expert_ids_(m / arg.sorted_tile_size_); @@ -138,30 +146,50 @@ struct ReferenceMoeGemm : public device::BaseOperator else { arg.b_element_op_(v_b, arg.b_e_n_k_(e, k, n)); - arg.b_element_op_(v_b_up, arg.b_e_n_k_(e, k, n + full_n)); + if constexpr(ActivationType == 2) + { + arg.b_element_op_(v_b_up, arg.b_e_n_k_(e, k, n + full_n)); + } } v_acc += ck::type_convert(v_a) * ck::type_convert(v_b); - v_acc_up += - ck::type_convert(v_a) * ck::type_convert(v_b_up); + + if constexpr(ActivationType == 2) + { + v_acc_up += ck::type_convert(v_a) * + ck::type_convert(v_b_up); + } } CDataType v_c{0}; CDataType v_c_up{0}; arg.c_element_op_(v_c, v_acc); - arg.c_element_op_(v_c_up, v_acc_up); - v_c = v_c * arg.b_scale_e_n_(e, n) * arg.a_scale_t_(t); - v_c = v_c * (1.0 / (1.0 + math::exp(-v_c))); - v_c_up = v_c_up * arg.b_scale_e_n_(e, n + full_n) * arg.a_scale_t_(t); - arg.c_t_k_n_(t, topk_id, n) = v_c * v_c_up; - // arg.c_t_k_n_(t, topk_id, n) = v_c + v_c_up; + if constexpr(ActivationType == 2) + { + arg.c_element_op_(v_c_up, v_acc_up); + v_c = v_c * arg.b_scale_e_n_(e, n) * arg.a_scale_t_(t); + v_c = v_c * (1.0 / (1.0 + math::exp(-v_c))); + v_c_up = v_c_up * arg.b_scale_e_n_(e, n + full_n) * arg.a_scale_t_(t); + arg.c_t_k_n_(t, topk_id, n) = v_c * v_c_up; + } + else + { + if constexpr(ActivationType == 1) + { + tensor_operation::element_wise::Silu{}(v_c, v_c); + } + else if constexpr(ActivationType == 0) + { + tensor_operation::element_wise::Gelu{}(v_c, v_c); + } + arg.c_t_k_n_(t, topk_id, n) = v_c; + } } }; const ck::index_t max_token_id = arg.max_token_id_(0); - make_ParallelTensorFunctor( - f_mk_kn_mn, max_token_id, full_n)( + make_ParallelTensorFunctor(f_mk_kn_mn, max_token_id, full_n)( std::thread::hardware_concurrency()); return 0;