19
csrc/moe/moe_gating_top_k_hash/CMakeLists.txt
Normal file
19
csrc/moe/moe_gating_top_k_hash/CMakeLists.txt
Normal file
@@ -0,0 +1,19 @@
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
# CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
# Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
# See LICENSE in the root of the software repository for the full text of the License.
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
|
||||
file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
|
||||
if(NOT ENABLE_TEST AND NOT BENCHMARK)
|
||||
list(REMOVE_ITEM CURRENT_DIRS tests)
|
||||
endif()
|
||||
foreach(SUB_DIR ${CURRENT_DIRS})
|
||||
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
|
||||
add_subdirectory(${SUB_DIR})
|
||||
endif()
|
||||
endforeach()
|
||||
62
csrc/moe/moe_gating_top_k_hash/op_host/CMakeLists.txt
Normal file
62
csrc/moe/moe_gating_top_k_hash/op_host/CMakeLists.txt
Normal file
@@ -0,0 +1,62 @@
|
||||
# This program is free software, you can redistribute it and/or modify it.
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
# This file is a part of the CANN Open Software.
|
||||
# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
# Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
# See LICENSE in the root of the software repository for the full text of the License.
|
||||
# ======================================================================================================================
|
||||
|
||||
# add_ops_compile_options(
|
||||
# OP_NAME MoeGatingTopKHash
|
||||
# OPTIONS --cce-auto-sync=off
|
||||
# -Wno-deprecated-declarations
|
||||
# -Werror
|
||||
# -mllvm -cce-aicore-hoist-movemask=false
|
||||
# --op_relocatable_kernel_binary=true
|
||||
# )
|
||||
|
||||
# set(moe_gating_top_k_hash_depends transformer/moe/moe_gating_top_k_hash PARENT_SCOPE)
|
||||
|
||||
# target_sources(op_host_aclnn PRIVATE
|
||||
# op_host/moe_gating_top_k_hash_def.cpp
|
||||
# )
|
||||
|
||||
# target_sources(optiling PRIVATE
|
||||
# op_host/moe_gating_top_k_hash_tiling.cpp
|
||||
# )
|
||||
|
||||
# if (NOT BUILD_OPEN_PROJECT)
|
||||
# target_sources(opmaster_ct PRIVATE
|
||||
# op_host/moe_gating_top_k_hash_tiling.cpp
|
||||
# )
|
||||
# endif ()
|
||||
|
||||
# target_include_directories(optiling PRIVATE
|
||||
# ${CMAKE_CURRENT_SOURCE_DIR}/op_host
|
||||
# )
|
||||
|
||||
# target_sources(opsproto PRIVATE
|
||||
# op_host/moe_gating_top_k_hash_proto.cpp
|
||||
# )
|
||||
|
||||
|
||||
if (BUILD_OPEN_PROJECT)
|
||||
message(STATUS "=== Debug: BUILD is BUILD_OPEN_PROJECT")
|
||||
target_sources(op_host_aclnn PRIVATE
|
||||
moe_gating_top_k_hash_def.cpp
|
||||
)
|
||||
add_ops_compile_options(
|
||||
OP_NAME MoeGatingTopKHash
|
||||
OPTIONS --cce-auto-sync=off
|
||||
-Wno-deprecated-declarations
|
||||
-mllvm -cce-aicore-hoist-movemask=false
|
||||
--op_relocatable_kernel_binary=true
|
||||
)
|
||||
endif()
|
||||
|
||||
if(NOT BUILD_OPS_RTY_KERNEL)
|
||||
message(STATUS "=== Debug: BUILD is NOT BUILD_OPS_RTY_KERNEL")
|
||||
add_modules_sources(OPTYPE moe_gating_top_k_hash ACLNNTYPE aclnn)
|
||||
endif()
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file moe_gating_top_k_hash_def.cpp
|
||||
* \brief
|
||||
*/
|
||||
#include "register/op_def_registry.h"
|
||||
|
||||
namespace ops {
|
||||
class MoeGatingTopKHash : public OpDef {
|
||||
public:
|
||||
explicit MoeGatingTopKHash(const char *name) : OpDef(name)
|
||||
{
|
||||
this->Input("x")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16,
|
||||
ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16,
|
||||
ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16,
|
||||
ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("bias")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16,
|
||||
ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16,
|
||||
ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16,
|
||||
ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("input_ids")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64,
|
||||
ge::DT_INT32, ge::DT_INT32, ge::DT_INT32,
|
||||
ge::DT_INT64, ge::DT_INT64, ge::DT_INT64,
|
||||
ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("tid2eid")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32,
|
||||
ge::DT_INT32, ge::DT_INT32, ge::DT_INT32,
|
||||
ge::DT_INT64, ge::DT_INT64, ge::DT_INT64,
|
||||
ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Output("y")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16,
|
||||
ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16,
|
||||
ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16,
|
||||
ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Output("expert_idx")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32,
|
||||
ge::DT_INT32, ge::DT_INT32, ge::DT_INT32,
|
||||
ge::DT_INT32, ge::DT_INT32, ge::DT_INT32,
|
||||
ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Output("out")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT,
|
||||
ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT,
|
||||
ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT,
|
||||
ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND,
|
||||
ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Attr("k").Int();
|
||||
this->Attr("k_group").AttrType(OPTIONAL).Int(1);
|
||||
this->Attr("group_count").AttrType(OPTIONAL).Int(1);
|
||||
this->Attr("group_select_mode").AttrType(OPTIONAL).Int(0);
|
||||
this->Attr("renorm").AttrType(OPTIONAL).Int(0);
|
||||
this->Attr("norm_type").AttrType(OPTIONAL).Int(0);
|
||||
this->Attr("out_flag").AttrType(OPTIONAL).Bool(false);
|
||||
this->Attr("routed_scaling_factor").AttrType(OPTIONAL).Float(1.0);
|
||||
this->Attr("eps").AttrType(OPTIONAL).Float(1e-20f);
|
||||
this->AICore().AddConfig("ascend910b");
|
||||
this->AICore().AddConfig("ascend910_93");
|
||||
|
||||
OpAICoreConfig regbaseCfg;
|
||||
regbaseCfg.DynamicCompileStaticFlag(true)
|
||||
.DynamicRankSupportFlag(true)
|
||||
.DynamicShapeSupportFlag(true)
|
||||
.ExtendCfgInfo("opFile.value", "moe_gating_top_k_hash");
|
||||
this->AICore().AddConfig("ascend950", regbaseCfg);
|
||||
}
|
||||
};
|
||||
|
||||
OP_ADD(MoeGatingTopKHash);
|
||||
} // namespace ops
|
||||
@@ -0,0 +1,68 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file moe_gating_top_k_hash_proto.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef OPS_OP_PROTO_INC_MOEGATINGTOPK_H_
|
||||
#define OPS_OP_PROTO_INC_MOEGATINGTOPK_H_
|
||||
|
||||
#include "graph/operator_reg.h"
|
||||
|
||||
namespace ge {
|
||||
|
||||
/**
|
||||
* @brief Compute renorm(sigmoid) and topk for moe input.
|
||||
*
|
||||
* @par Inputs:
|
||||
* @li x: A 2D tensor which moe gating topk is applied, The shape is: (B*S, E), format supports ND, and data type must be float16, float or bfloat16. E(Expert num) can not be greater than 2048. E(Expert num) should be divisible by group_count.
|
||||
* @li bias: A 1D tensor which is "bias" in moe gating topk. The shape is: (E), format supports ND, and data type must be the same as that of x.
|
||||
*
|
||||
* @par Outputs:
|
||||
* @li y: A 2D tensor which is the topk value result of moe gating topk, format supports ND, and data type must be the same as that of x.
|
||||
The size of the non-1 axis must be the same as that of the corresponding axis of x.
|
||||
The size of the -1 axis must be the same as that of k.
|
||||
* @li expert_idx: A 2D tensor which is the topk index result of moe gating topk, format supports ND, and data type must be int. The shape must be the same as that of y.
|
||||
* @li out: A 2D tensor which is the renorm result of moe gating topk, format supports ND, and data type must be float. The shape must be the same as that of x.
|
||||
*
|
||||
* @par Attributes:
|
||||
* @li k: A required attribute of type int. The value must greater than 0 and less than or equal to expert_num / group_count * k_group, indicating the topk value.
|
||||
* @li k_group: An optional attribute of type int. It can not be less than 1, and can not be greater than group_count, indicating the topk group value. The default value is 1.
|
||||
* @li group_count: An optional attribute of type int. It can not be less than 1, indicating the group count. The group_count * align_32(expert_num / group_count) can not be greater than 2048. The default value is 1.
|
||||
* @li group_select_mode: An optional attribute of type int. 0 indicating that sort group by max values, 1 indicating that sort group by sum of top-2 values. The default value is 0.
|
||||
* @li renorm: An optional attribute of type int. It can only be 0 now, indicating that norm firstly and then topk. The default value is 0.
|
||||
* @li norm_type: An optional attribute of type int. 0 indicating that the softmax function is used, 1 indicating that the sigmoid function is used. The default value is 0.
|
||||
* @li out_flag: An optional attribute of type bool. true indicating that has renorm output, false indicating that does not have renorm output. The default value is false.
|
||||
* @li routed_scaling_factor: An optional attribute of type float, indicating the routed_scaling_factor coefficient in use. The default value is 1.0.
|
||||
* @li eps: An optional attribute of type float, indicating the eps coefficient in use. The default value is 1e-20.
|
||||
*/
|
||||
REG_OP(MoeGatingTopKHash)
|
||||
.INPUT(x, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16}))
|
||||
.OPTIONAL_INPUT(bias, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16}))
|
||||
.OPTIONAL_INPUT(input_ids, TensorType({DT_INT64, DT_INT32}))
|
||||
.OPTIONAL_INPUT(tid2eid, TensorType({DT_INT64, DT_INT32}))
|
||||
.OUTPUT(y, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16}))
|
||||
.OUTPUT(expert_idx, TensorType({DT_INT32}))
|
||||
.OUTPUT(out, TensorType({DT_FLOAT}))
|
||||
.REQUIRED_ATTR(k, Int)
|
||||
.ATTR(k_group, Int, 1)
|
||||
.ATTR(group_count, Int, 1)
|
||||
.ATTR(group_select_mode, Int, 0)
|
||||
.ATTR(renorm, Int, 0)
|
||||
.ATTR(norm_type, Int, 0)
|
||||
.ATTR(out_flag, Bool, false)
|
||||
.ATTR(routed_scaling_factor, Float, 1.0)
|
||||
.ATTR(eps, Float, 1e-20f)
|
||||
.OP_END_FACTORY_REG(MoeGatingTopKHash)
|
||||
|
||||
} // namespace ge
|
||||
|
||||
#endif // OPS_OP_PROTO_INC_MOEGATINGTOPK_H_
|
||||
@@ -0,0 +1,640 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/* !
|
||||
* \file moe_gating_top_k_hash_tiling.cpp
|
||||
* \brief
|
||||
*/
|
||||
#include "moe_gating_top_k_hash_tiling.h"
|
||||
#include "moe_gating_top_k_hash_tiling_arch35.h"
|
||||
#include <graph/utils/type_utils.h>
|
||||
|
||||
namespace optiling {
|
||||
const static int64_t GROUP_SELECT_MODE_MAX = 0;
|
||||
const static int64_t GROUP_SELECT_MODE_SUM = 1;
|
||||
const static int64_t RENORM_NO = 0;
|
||||
const static int64_t RENORM_L1 = 1;
|
||||
const static int64_t NORM_TYPE_SOFTMAX = 0;
|
||||
const static int64_t NORM_TYPE_SIGMOID = 1;
|
||||
const static int64_t NORM_TYPE_SOFTPLUS = 2;
|
||||
const static int64_t OUT_FLAG_FALSE = 0;
|
||||
const static int64_t OUT_FLAG_TRUE = 1;
|
||||
const static size_t X_INPUT_DIMS = 2;
|
||||
const static size_t BIAS_INPUT_DIMS = 1;
|
||||
const static size_t Y_OUTPUT_DIMS = 2;
|
||||
const static size_t EXPERT_IDX_OUTPUY_DIMS = 2;
|
||||
const static size_t OUT_OUTPUT_DIMS = 2;
|
||||
const static int64_t MAX_EXPERT_COUNT = 2048;
|
||||
|
||||
const static int64_t X_INPUT_INDEX = 0;
|
||||
const static int64_t BIAS_INPUT_INDEX = 1;
|
||||
const static int64_t INPUT_IDS_INPUT_INDEX = 2;
|
||||
const static int64_t TID_TO_EID_INPUT_INDEX = 3;
|
||||
const static int64_t Y_OUTPUT_INDEX = 0;
|
||||
const static int64_t EXPERT_IDX_OUTPUT_INDEX = 1;
|
||||
const static int64_t OUT_OUTPUT_INDEX = 2;
|
||||
const static int64_t K_ATTR_INDEX = 0;
|
||||
const static int64_t K_GROUP_ATTR_INDEX = 1;
|
||||
const static int64_t GROUP_COUNT_ATTR_INDEX = 2;
|
||||
const static int64_t GROUP_SELECT_MODE_ATTR_INDEX = 3;
|
||||
const static int64_t RENORM_ATTR_INDEX = 4;
|
||||
const static int64_t NORM_TYPE_ATTR_INDEX = 5;
|
||||
const static int64_t OUT_FLAG_ATTR_INDEX = 6;
|
||||
const static int64_t ROUTED_SCALING_FACTOR_ATTR_INDEX = 7;
|
||||
const static int64_t EPS_ATTR_INDEX = 8;
|
||||
const static int64_t DEFAULT_WORKSPACE_SIZE = 16777216; // 预留16M空间
|
||||
const static uint32_t DATATYPESIZE_FLOAT = 4;
|
||||
const static bool IS_LARGEST = true;
|
||||
const static bool IS_INITINDEX = false;
|
||||
const static bool IS_REUSESOURCE = false;
|
||||
const static uint64_t WITH_GROUP_CONDITION = 1;
|
||||
const static uint64_t WITHOUT_GROUP_CONDITION = 2;
|
||||
const static uint64_t MAX_IN_GROUP_CONDITION = 3;
|
||||
constexpr int32_t ROW_COUNT_PER_TASK = 1;
|
||||
|
||||
const static uint64_t TILING_KEY_EXPERTNUM_GROUPNUM_ALIGN_HIGH_PERF = 0;
|
||||
const static uint64_t TILING_KEY_WITHOUT_GROUP = 1;
|
||||
const static uint64_t TILING_KEY_GENERALIZED = 2;
|
||||
const static uint64_t TILING_KEY_WITHOUT_GROUP_0 = 3;
|
||||
const static uint64_t TILING_KEY_WITHOUT_GROUP_1 = 4;
|
||||
const static uint64_t TILING_KEY_WITHOUT_GROUP_2 = 5;
|
||||
const static uint64_t TILING_KEY_WITHOUT_GROUP_3 = 6;
|
||||
|
||||
template <typename T>
|
||||
static inline T CeilDiv(T num, T rnd)
|
||||
{
|
||||
return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)));
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
static inline T CeilAlign(T num, T rnd)
|
||||
{
|
||||
return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)) * (rnd));
|
||||
}
|
||||
|
||||
inline static int64_t CeilLog4(int64_t x)
|
||||
{
|
||||
return static_cast<int64_t>(std::ceil(std::log(x) / std::log(4))); // 4 for four
|
||||
}
|
||||
|
||||
class MoeGatingTopKHashTilingBase {
|
||||
public:
|
||||
explicit MoeGatingTopKHashTilingBase(gert::TilingContext *context) : context_(context)
|
||||
{
|
||||
Reset();
|
||||
}
|
||||
~MoeGatingTopKHashTilingBase() = default;
|
||||
|
||||
bool IsCapable()
|
||||
{
|
||||
return true;
|
||||
}
|
||||
// 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小
|
||||
ge::graphStatus GetPlatformInfo() ;
|
||||
// 2、获取INPUT/OUTPUT/ATTR信息
|
||||
ge::graphStatus GetShapeAttrsInfo() ;
|
||||
// 3、计算数据切分TilingData
|
||||
ge::graphStatus DoOpTiling() ;
|
||||
// 4、计算高阶API的TilingData
|
||||
ge::graphStatus DoLibApiTiling() ;
|
||||
// 5、计算TilingKey
|
||||
uint64_t GetTilingKey() const ;
|
||||
// 6、计算Workspace 大小
|
||||
ge::graphStatus GetWorkspaceSize() ;
|
||||
// 7、保存Tiling数据
|
||||
ge::graphStatus PostTiling() ;
|
||||
void Reset();
|
||||
|
||||
private:
|
||||
ge::graphStatus CheckInputShape();
|
||||
ge::graphStatus CheckAttr();
|
||||
ge::graphStatus CheckOutShape();
|
||||
void SplitRows();
|
||||
void CalTmpBufUbSize();
|
||||
|
||||
const gert::Shape *xShape_ = nullptr;
|
||||
const gert::Shape *biasShape_ = nullptr;
|
||||
const gert::Shape *inputIdsShape_ = nullptr;
|
||||
const gert::Shape *tid2eidShape_ = nullptr;
|
||||
const gert::Shape *yShape_ = nullptr;
|
||||
const gert::Shape *expertIdxShape_ = nullptr;
|
||||
const gert::Shape *outShape_ = nullptr;
|
||||
ge::DataType inputIdsDtype;
|
||||
ge::DataType tid2eidDtype;
|
||||
|
||||
uint64_t coreNum_ = 0;
|
||||
int64_t rows_ = 0;
|
||||
int64_t expertCount_ = 0;
|
||||
int64_t addBias_ = 0;
|
||||
|
||||
int64_t k_ = 0;
|
||||
int64_t kGroup_ = 0;
|
||||
int64_t groupCount_ = 0;
|
||||
int64_t perGroupExpertCount_ = 0;
|
||||
int64_t groupSelectMode_ = GROUP_SELECT_MODE_MAX;
|
||||
int64_t renorm_ = RENORM_NO;
|
||||
int64_t normType_ = NORM_TYPE_SOFTMAX;
|
||||
int64_t outFlag_ = OUT_FLAG_FALSE;
|
||||
int64_t hashFlag_ = 0;
|
||||
float routedScalingFactor_ = 1.0;
|
||||
float eps_ = 1e-20f;
|
||||
|
||||
int64_t inputDtypeSize_;
|
||||
const char *opName_ = "";
|
||||
MoeGatingTopKHashTilingData moeGatingTopKTilingData_;
|
||||
gert::TilingContext *context_ = nullptr;
|
||||
uint64_t workspaceSize_ = 0;
|
||||
};
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingBase::CheckInputShape()
|
||||
{
|
||||
size_t xDimNum = xShape_->GetDimNum();
|
||||
OPS_ERR_IF(xDimNum != X_INPUT_DIMS,
|
||||
OPS_LOG_E(context_, "The dim number of x is: %zu, but should be %zu.", xDimNum, X_INPUT_DIMS),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
// 通过输入获取rows 和 expertCount
|
||||
rows_ = xShape_->GetDim(0);
|
||||
expertCount_ = xShape_->GetDim(1);
|
||||
moeGatingTopKTilingData_.set_rowCount(rows_);
|
||||
moeGatingTopKTilingData_.set_expertCount(expertCount_);
|
||||
if (biasShape_ != nullptr) {
|
||||
addBias_ = 1;
|
||||
size_t biasDimNum = biasShape_->GetDimNum();
|
||||
OPS_ERR_IF(biasDimNum != BIAS_INPUT_DIMS,
|
||||
OPS_LOG_E(context_, "The dim number of bias is: %zu, but should be %zu.", biasDimNum, BIAS_INPUT_DIMS),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF(
|
||||
biasShape_->GetDim(0) != expertCount_,
|
||||
OPS_LOG_E(context_, "The first dim of bias is: %ld, but should be %ld.", biasShape_->GetDim(0), expertCount_),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
moeGatingTopKTilingData_.set_addBias(addBias_);
|
||||
|
||||
if (inputIdsShape_ != nullptr) {
|
||||
OPS_ERR_IF(
|
||||
tid2eidShape_ == nullptr,
|
||||
OPS_LOG_E(context_, "The tid2eid should not be empty when inputIds has value."),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
if (tid2eidShape_ != nullptr) {
|
||||
OPS_ERR_IF(
|
||||
inputIdsShape_ == nullptr,
|
||||
OPS_LOG_E(context_, "The inputIds should not be empty when tid2eid has value."),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
if (inputIdsShape_ != nullptr && tid2eidShape_ != nullptr) {
|
||||
hashFlag_ = 1;
|
||||
OPS_LOG_I(context_, "hashFlag_ is 1.");
|
||||
}
|
||||
moeGatingTopKTilingData_.set_hashFlag(hashFlag_);
|
||||
|
||||
OPS_ERR_IF(k_ > expertCount_,
|
||||
OPS_LOG_E(context_, "k is: %ld, expert num is: %ld, k cannot be greater than expert num.", k_, expertCount_),
|
||||
return ge::GRAPH_FAILED);
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingBase::CheckAttr()
|
||||
{
|
||||
OPS_ERR_IF(
|
||||
expertCount_ > MAX_EXPERT_COUNT,
|
||||
OPS_LOG_E(context_, "expert count is: %ld, but should not greater than %ld.", expertCount_, MAX_EXPERT_COUNT),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
OPS_ERR_IF(k_ <= 0, OPS_LOG_E(context_, "k is: %ld, but should be greater than 0.", k_), return ge::GRAPH_FAILED);
|
||||
|
||||
OPS_ERR_IF(kGroup_ <= 0, OPS_LOG_E(context_, "k_group is: %ld, but should be greater than 0.", kGroup_),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
OPS_ERR_IF(kGroup_ > groupCount_,
|
||||
OPS_LOG_E(context_, "k_group is: %ld, but should not greater than %ld.", kGroup_, groupCount_),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
OPS_ERR_IF(groupCount_ <= 0, OPS_LOG_E(context_, "group_count is: %ld, but should be greater than 0.", groupCount_),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
OPS_ERR_IF(normType_ != NORM_TYPE_SOFTMAX && normType_ != NORM_TYPE_SIGMOID && normType_ != NORM_TYPE_SOFTPLUS,
|
||||
OPS_LOG_E(context_, "norm type is: %ld, but currently only support %ld, %ld and %ld.", normType_,
|
||||
NORM_TYPE_SOFTMAX, NORM_TYPE_SIGMOID, NORM_TYPE_SOFTPLUS),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
OPS_ERR_IF(normType_ == NORM_TYPE_SOFTPLUS && groupCount_ != 1,
|
||||
OPS_LOG_E(context_, "norm type softplus only supported when groupCount equals 1, but got %ld.", groupCount_),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
OPS_ERR_IF(groupSelectMode_ != GROUP_SELECT_MODE_SUM && groupSelectMode_ != GROUP_SELECT_MODE_MAX,
|
||||
OPS_LOG_E(context_, "group select mode is: %ld, but currently only support %ld and %ld.", groupSelectMode_,
|
||||
GROUP_SELECT_MODE_SUM, GROUP_SELECT_MODE_MAX),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
OPS_ERR_IF(renorm_ != RENORM_NO,
|
||||
OPS_LOG_E(context_, "renorm is: %ld, but currently only support %ld.", renorm_, RENORM_NO),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
OPS_ERR_IF(expertCount_ % groupCount_ != 0,
|
||||
OPS_LOG_E(context_, "Expert count : %ld is not divisible by k_group: %ld", expertCount_, groupCount_),
|
||||
return ge::GRAPH_FAILED);
|
||||
perGroupExpertCount_ = expertCount_ / groupCount_;
|
||||
|
||||
OPS_ERR_IF(perGroupExpertCount_ < 1,
|
||||
OPS_LOG_E(context_, "group expert count is: %ld, but should be greater than 1.", perGroupExpertCount_),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF(
|
||||
groupSelectMode_ == GROUP_SELECT_MODE_SUM && perGroupExpertCount_ < 2,
|
||||
OPS_LOG_E(context_,
|
||||
"group expert count is: %ld, if group select mode is: %ld, group expert count should be greater than 1.",
|
||||
perGroupExpertCount_, groupSelectMode_),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF(k_ > kGroup_ * perGroupExpertCount_,
|
||||
OPS_LOG_E(context_, "k is: %ld, but should be smaller than %ld.", k_, kGroup_ * perGroupExpertCount_),
|
||||
return ge::GRAPH_FAILED);
|
||||
int64_t groupExpertCountAlign = CeilAlign(perGroupExpertCount_, 32L);
|
||||
if (groupCount_ != 1 && groupCount_ != expertCount_ && kGroup_ != groupCount_) {
|
||||
// 分组场景下才需要校验对齐后的数量
|
||||
OPS_ERR_IF(groupCount_ * groupExpertCountAlign > MAX_EXPERT_COUNT,
|
||||
OPS_LOG_E(context_, "group count * group expert count align is: %ld, but should not greater than %ld.",
|
||||
groupCount_ * groupExpertCountAlign, MAX_EXPERT_COUNT),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
|
||||
moeGatingTopKTilingData_.set_perGroupExpertCount(perGroupExpertCount_);
|
||||
moeGatingTopKTilingData_.set_perGroupExpertCountAlign(groupExpertCountAlign);
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingBase::GetShapeAttrsInfo()
|
||||
{
|
||||
opName_ = context_->GetNodeName();
|
||||
// 获取输入shape信息
|
||||
auto xShapePtr = context_->GetInputShape(X_INPUT_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, xShapePtr, return ge::GRAPH_FAILED);
|
||||
xShape_ = &xShapePtr->GetStorageShape();
|
||||
|
||||
auto biasShapePtr = context_->GetOptionalInputShape(BIAS_INPUT_INDEX);
|
||||
biasShape_ = biasShapePtr == nullptr ? nullptr : &biasShapePtr->GetStorageShape();
|
||||
|
||||
auto inputIdsShapePtr = context_->GetOptionalInputShape(INPUT_IDS_INPUT_INDEX);
|
||||
inputIdsShape_ = inputIdsShapePtr == nullptr ? nullptr : &inputIdsShapePtr->GetStorageShape();
|
||||
|
||||
auto tid2eidShapePtr = context_->GetOptionalInputShape(TID_TO_EID_INPUT_INDEX);
|
||||
tid2eidShape_ = tid2eidShapePtr == nullptr ? nullptr : &tid2eidShapePtr->GetStorageShape();
|
||||
|
||||
// 获取输出shape
|
||||
auto yShapePtr = context_->GetOutputShape(Y_OUTPUT_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, yShapePtr, return ge::GRAPH_FAILED);
|
||||
yShape_ = &yShapePtr->GetStorageShape();
|
||||
auto expertIdxPtr = context_->GetOutputShape(EXPERT_IDX_OUTPUT_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, expertIdxPtr, return ge::GRAPH_FAILED);
|
||||
expertIdxShape_ = &expertIdxPtr->GetStorageShape();
|
||||
auto outPtr = context_->GetOutputShape(OUT_OUTPUT_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, outPtr, return ge::GRAPH_FAILED);
|
||||
outShape_ = &outPtr->GetStorageShape();
|
||||
|
||||
auto x = context_->GetInputDesc(X_INPUT_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, x, return ge::GRAPH_FAILED);
|
||||
auto xDtype = x->GetDataType();
|
||||
OPS_ERR_IF(
|
||||
(xDtype != ge::DataType::DT_FLOAT && xDtype != ge::DataType::DT_FLOAT16 && xDtype != ge::DataType::DT_BF16),
|
||||
OPS_LOG_E(context_, "x dtype %s error, only supports float32, half, bf16. please check.",
|
||||
ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
if (biasShapePtr != nullptr) {
|
||||
auto biasDtype = context_->GetOptionalInputDesc(BIAS_INPUT_INDEX)->GetDataType();
|
||||
OPS_ERR_IF((biasDtype != xDtype),
|
||||
OPS_LOG_E(context_, "bias dtype %s not equal x dtype %s, please check.",
|
||||
ge::TypeUtils::DataTypeToSerialString(biasDtype).c_str(),
|
||||
ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
if (inputIdsShapePtr != nullptr) {
|
||||
inputIdsDtype = context_->GetOptionalInputDesc(INPUT_IDS_INPUT_INDEX)->GetDataType();
|
||||
OPS_ERR_IF((inputIdsDtype != ge::DataType::DT_INT32 && inputIdsDtype != ge::DataType::DT_INT64),
|
||||
OPS_LOG_E(context_, "inputIds dtype %s error, only supports int32 and int64. please check.",
|
||||
ge::TypeUtils::DataTypeToSerialString(inputIdsDtype).c_str()),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
if (tid2eidShapePtr != nullptr) {
|
||||
tid2eidDtype = context_->GetOptionalInputDesc(TID_TO_EID_INPUT_INDEX)->GetDataType();
|
||||
OPS_ERR_IF((tid2eidDtype != ge::DataType::DT_INT32 && tid2eidDtype != ge::DataType::DT_INT64),
|
||||
OPS_LOG_E(context_, "tid2eid dtype %s error, only supports int32 and int64. please check.",
|
||||
ge::TypeUtils::DataTypeToSerialString(tid2eidDtype).c_str()),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
|
||||
auto yDesc = context_->GetOutputDesc(Y_OUTPUT_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, yDesc, return ge::GRAPH_FAILED);
|
||||
auto yDtype = yDesc->GetDataType();
|
||||
OPS_ERR_IF((yDtype != xDtype),
|
||||
OPS_LOG_E(context_, "y out dtype %s must be the same with x dtype %s.",
|
||||
ge::TypeUtils::DataTypeToSerialString(yDtype).c_str(),
|
||||
ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
auto expertIdDesc = context_->GetOutputDesc(EXPERT_IDX_OUTPUT_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, expertIdDesc, return ge::GRAPH_FAILED);
|
||||
auto expertIdDtype = expertIdDesc->GetDataType();
|
||||
OPS_ERR_IF((expertIdDtype != ge::DataType::DT_INT32),
|
||||
OPS_LOG_E(context_, "expertId out dtype %s error, only supports int32. please check.",
|
||||
ge::TypeUtils::DataTypeToSerialString(expertIdDtype).c_str()),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
auto normOutDesc = context_->GetOutputDesc(OUT_OUTPUT_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, normOutDesc, return ge::GRAPH_FAILED);
|
||||
auto normOutDtype = normOutDesc->GetDataType();
|
||||
OPS_ERR_IF((normOutDtype != ge::DataType::DT_FLOAT),
|
||||
OPS_LOG_E(context_, "norm out dtype %s error, only supports float. please check.",
|
||||
ge::TypeUtils::DataTypeToSerialString(normOutDtype).c_str()),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
// 获取属性
|
||||
auto attrs = context_->GetAttrs();
|
||||
OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED);
|
||||
|
||||
const int64_t *kPtr = attrs->GetAttrPointer<int64_t>(K_ATTR_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, kPtr, return ge::GRAPH_FAILED);
|
||||
k_ = *kPtr;
|
||||
moeGatingTopKTilingData_.set_k(k_);
|
||||
OPS_LOG_I(context_, "Attr k is: %ld ", k_);
|
||||
|
||||
const int64_t *kGroupPtr = attrs->GetAttrPointer<int64_t>(K_GROUP_ATTR_INDEX);
|
||||
if (kGroupPtr != nullptr) {
|
||||
kGroup_ = *kGroupPtr;
|
||||
moeGatingTopKTilingData_.set_kGroup(kGroup_);
|
||||
}
|
||||
OPS_LOG_I(context_, "Attr k_group is: %ld ", kGroup_);
|
||||
|
||||
const int64_t *groupCountPtr = attrs->GetAttrPointer<int64_t>(GROUP_COUNT_ATTR_INDEX);
|
||||
if (groupCountPtr != nullptr) {
|
||||
groupCount_ = *groupCountPtr;
|
||||
moeGatingTopKTilingData_.set_groupCount(groupCount_);
|
||||
}
|
||||
OPS_LOG_I(context_, "Attr group_count is: %ld ", groupCount_);
|
||||
|
||||
const int64_t *groupSelectModePtr = attrs->GetAttrPointer<int64_t>(GROUP_SELECT_MODE_ATTR_INDEX);
|
||||
if (groupSelectModePtr != nullptr) {
|
||||
groupSelectMode_ = *groupSelectModePtr;
|
||||
moeGatingTopKTilingData_.set_groupSelectMode(groupSelectMode_);
|
||||
}
|
||||
OPS_LOG_I(context_, "Attr group_select_mode is: %ld ", groupSelectMode_);
|
||||
|
||||
const int64_t *renormPtr = attrs->GetAttrPointer<int64_t>(RENORM_ATTR_INDEX);
|
||||
if (renormPtr != nullptr) {
|
||||
renorm_ = *renormPtr;
|
||||
moeGatingTopKTilingData_.set_renorm(renorm_);
|
||||
}
|
||||
OPS_LOG_I(context_, "Attr renorm is: %ld ", renorm_);
|
||||
|
||||
const int64_t *normTypePtr = attrs->GetAttrPointer<int64_t>(NORM_TYPE_ATTR_INDEX);
|
||||
if (normTypePtr != nullptr) {
|
||||
normType_ = *normTypePtr;
|
||||
moeGatingTopKTilingData_.set_normType(normType_);
|
||||
}
|
||||
OPS_LOG_I(context_, "Attr norm_type is: %ld ", normType_);
|
||||
|
||||
const bool *outFlagPtr = attrs->GetAttrPointer<bool>(OUT_FLAG_ATTR_INDEX);
|
||||
if (outFlagPtr != nullptr) {
|
||||
outFlag_ = (*outFlagPtr) ? 1 : 0;
|
||||
moeGatingTopKTilingData_.set_outFlag(outFlag_);
|
||||
}
|
||||
OPS_LOG_I(context_, "Attr out_flag is: %ld ", outFlag_);
|
||||
|
||||
const float *routedScalingFactorPtr = attrs->GetAttrPointer<float>(ROUTED_SCALING_FACTOR_ATTR_INDEX);
|
||||
if (routedScalingFactorPtr != nullptr) {
|
||||
routedScalingFactor_ = *routedScalingFactorPtr;
|
||||
moeGatingTopKTilingData_.set_routedScalingFactor(routedScalingFactor_);
|
||||
}
|
||||
OPS_LOG_I(context_, "Attr routed_scaling_factor is: %f ", routedScalingFactor_);
|
||||
|
||||
const float *epsPtr = attrs->GetAttrPointer<float>(EPS_ATTR_INDEX);
|
||||
if (epsPtr != nullptr) {
|
||||
eps_ = *epsPtr;
|
||||
moeGatingTopKTilingData_.set_eps(eps_);
|
||||
}
|
||||
OPS_LOG_I(context_, "Attr eps is: %f ", eps_);
|
||||
|
||||
inputDtypeSize_ = static_cast<int64_t>(ge::GetSizeByDataType(context_->GetInputDesc(0)->GetDataType()));
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingBase::GetPlatformInfo()
|
||||
{
|
||||
auto platformInfo = context_->GetPlatformInfo();
|
||||
OPS_ERR_IF(platformInfo == nullptr, OPS_LOG_E(context_, "fail to get platform info"), return ge::GRAPH_FAILED);
|
||||
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
|
||||
coreNum_ = ascendcPlatform.GetCoreNumAiv();
|
||||
uint64_t ubSizePlatForm;
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingBase::CheckOutShape()
|
||||
{
|
||||
OPS_ERR_IF((yShape_->GetDimNum() != xShape_->GetDimNum()),
|
||||
OPS_LOG_E(context_, "y out shape num %zu and x shape num %zu not equal, please check.", yShape_->GetDimNum(),
|
||||
xShape_->GetDimNum()),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF((expertIdxShape_->GetDimNum() != xShape_->GetDimNum()),
|
||||
OPS_LOG_E(context_, "expertId out shape num %zu and x shape num %zu not equal, please check.",
|
||||
expertIdxShape_->GetDimNum(), xShape_->GetDimNum()),
|
||||
return ge::GRAPH_FAILED);
|
||||
if (outShape_ != nullptr) {
|
||||
OPS_ERR_IF((outShape_->GetDimNum() != xShape_->GetDimNum()),
|
||||
OPS_LOG_E(context_, "norm out shape num %zu and x shape num %zu not equal, please check.",
|
||||
outShape_->GetDimNum(), xShape_->GetDimNum()),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
|
||||
OPS_ERR_IF((yShape_->GetDim(0) != xShape_->GetDim(0)),
|
||||
OPS_LOG_E(context_, "y out dim[0] %ld not equal x dim[0] %ld, please check.", yShape_->GetDim(0),
|
||||
xShape_->GetDim(0)),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF((expertIdxShape_->GetDim(0) != xShape_->GetDim(0)),
|
||||
OPS_LOG_E(context_, "expertId out dim[0] %ld not equal x dim[0] %ld, please check.",
|
||||
expertIdxShape_->GetDim(0), xShape_->GetDim(0)),
|
||||
return ge::GRAPH_FAILED);
|
||||
if (outFlag_ && outShape_ != nullptr) {
|
||||
OPS_ERR_IF((outShape_->GetDim(0) != xShape_->GetDim(0)),
|
||||
OPS_LOG_E(context_, "norm out dim[0] %ld and x dim[0] %ld not equal, please check.",
|
||||
outShape_->GetDim(0), outShape_->GetDim(0)),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
|
||||
OPS_ERR_IF((yShape_->GetDim(1) != k_),
|
||||
OPS_LOG_E(context_, "y dim[1] %ld not equal k %ld, please check.", yShape_->GetDim(1), k_),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF((expertIdxShape_->GetDim(1) != k_),
|
||||
OPS_LOG_E(context_, "expertId dim[1] %ld not equal k %ld, please check.", expertIdxShape_->GetDim(1), k_),
|
||||
return ge::GRAPH_FAILED);
|
||||
if (outFlag_ && outShape_ != nullptr) {
|
||||
OPS_ERR_IF((outShape_->GetDim(1) != xShape_->GetDim(1)),
|
||||
OPS_LOG_E(context_, "normOut dim[1] %ld and x dim[1] %ld not equal, please check.", outShape_->GetDim(1),
|
||||
xShape_->GetDim(1)),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
void MoeGatingTopKHashTilingBase::SplitRows()
|
||||
{
|
||||
int64_t perCoreRows = CeilDiv(rows_, static_cast<int64_t>(coreNum_));
|
||||
int64_t needCoreNum = CeilDiv(rows_, perCoreRows);
|
||||
// perCoreRows cannot be 0
|
||||
int64_t lastCoreRows = rows_ % perCoreRows == 0 ? perCoreRows : rows_ % perCoreRows;
|
||||
moeGatingTopKTilingData_.set_needCoreNum(needCoreNum);
|
||||
moeGatingTopKTilingData_.set_perCoreRowCount(perCoreRows);
|
||||
moeGatingTopKTilingData_.set_lastCoreRowCount(lastCoreRows);
|
||||
int64_t vmsCount = CeilLog4(CeilDiv(kGroup_, 4L));
|
||||
OPS_LOG_I(context_, "vms count is: %ld", vmsCount);
|
||||
moeGatingTopKTilingData_.set_vmsCount(vmsCount); // 需要归并的轮数
|
||||
}
|
||||
|
||||
void MoeGatingTopKHashTilingBase::CalTmpBufUbSize()
|
||||
{
|
||||
std::vector<int64_t> shape_vec = {expertCount_};
|
||||
ge::Shape shape(shape_vec);
|
||||
uint32_t maxValue = 0;
|
||||
uint32_t minValue = 0;
|
||||
AscendC::GetSigmoidMaxMinTmpSize(shape, sizeof(float), false, maxValue, minValue);
|
||||
|
||||
int64_t indexTmpBuf = (expertCount_ + 31) / 32 * 32 * static_cast<int64_t>(sizeof(float));
|
||||
moeGatingTopKTilingData_.set_calTmpBufUbSize(std::max(indexTmpBuf, static_cast<int64_t>(minValue)));
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingBase::DoOpTiling()
|
||||
{
|
||||
auto ret = GetPlatformInfo();
|
||||
if (ret != ge::GRAPH_SUCCESS) {
|
||||
return ret;
|
||||
}
|
||||
|
||||
ret = GetShapeAttrsInfo();
|
||||
if (ret != ge::GRAPH_SUCCESS) {
|
||||
return ret;
|
||||
}
|
||||
|
||||
ret = CheckInputShape();
|
||||
if (ret != ge::GRAPH_SUCCESS) {
|
||||
return ret;
|
||||
}
|
||||
|
||||
ret = CheckOutShape();
|
||||
if (ret != ge::GRAPH_SUCCESS) {
|
||||
return ret;
|
||||
}
|
||||
|
||||
ret = CheckAttr();
|
||||
if (ret != ge::GRAPH_SUCCESS) {
|
||||
return ret;
|
||||
}
|
||||
|
||||
CalTmpBufUbSize();
|
||||
SplitRows();
|
||||
|
||||
ret = PostTiling();
|
||||
if (ret != ge::GRAPH_SUCCESS) {
|
||||
return ret;
|
||||
}
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingBase::DoLibApiTiling()
|
||||
{
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingBase::GetWorkspaceSize()
|
||||
{
|
||||
// 计算workspace大小
|
||||
workspaceSize_ = DEFAULT_WORKSPACE_SIZE;
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingBase::PostTiling()
|
||||
{
|
||||
context_->SetTilingKey(GetTilingKey());
|
||||
context_->SetBlockDim(moeGatingTopKTilingData_.get_needCoreNum());
|
||||
size_t *currentWorkspace = context_->GetWorkspaceSizes(1);
|
||||
currentWorkspace[0] = workspaceSize_;
|
||||
moeGatingTopKTilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(),
|
||||
context_->GetRawTilingData()->GetCapacity());
|
||||
context_->GetRawTilingData()->SetDataSize(moeGatingTopKTilingData_.GetDataSize());
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
uint64_t MoeGatingTopKHashTilingBase::GetTilingKey() const
|
||||
{
|
||||
// DeepSeekV3排序对齐高性能场景
|
||||
if (expertCount_ == 256 && groupCount_ == 8 && kGroup_ == 4 && k_ <= 32 && addBias_ &&
|
||||
groupSelectMode_ == GROUP_SELECT_MODE_SUM && renorm_ == RENORM_NO && normType_ == NORM_TYPE_SIGMOID &&
|
||||
!outFlag_) {
|
||||
// DeepSeekV3排序对齐高性能场景
|
||||
return TILING_KEY_EXPERTNUM_GROUPNUM_ALIGN_HIGH_PERF;
|
||||
} else if (groupCount_ == 1 || groupCount_ == expertCount_ || kGroup_ == groupCount_) {
|
||||
/**
|
||||
* 不分组场景：
|
||||
* 1. 分组数为 1
|
||||
* 2. 分组数等于专家数（每个组只有一个专家）
|
||||
* 3. 选择所有组
|
||||
*/
|
||||
if (inputIdsShape_ == nullptr) {
|
||||
return TILING_KEY_WITHOUT_GROUP;
|
||||
} else if (inputIdsDtype == ge::DataType::DT_INT32 && tid2eidDtype == ge::DataType::DT_INT64) {
|
||||
return TILING_KEY_WITHOUT_GROUP_0;
|
||||
} else if (inputIdsDtype == ge::DataType::DT_INT32 && tid2eidDtype == ge::DataType::DT_INT32) {
|
||||
return TILING_KEY_WITHOUT_GROUP_1;
|
||||
} else if (inputIdsDtype == ge::DataType::DT_INT64 && tid2eidDtype == ge::DataType::DT_INT64) {
|
||||
return TILING_KEY_WITHOUT_GROUP_2;
|
||||
} else if (inputIdsDtype == ge::DataType::DT_INT64 && tid2eidDtype == ge::DataType::DT_INT32) {
|
||||
return TILING_KEY_WITHOUT_GROUP_3;
|
||||
}
|
||||
} else {
|
||||
return TILING_KEY_GENERALIZED;
|
||||
}
|
||||
}
|
||||
|
||||
void MoeGatingTopKHashTilingBase::Reset()
|
||||
{
|
||||
opName_ = nullptr;
|
||||
return;
|
||||
}
|
||||
|
||||
ge::graphStatus TilingForMoeGatingTopKHash(gert::TilingContext *context)
|
||||
{
|
||||
OPS_LOG_I(context, "TilingForMoeGatingTopKHash start");
|
||||
OPS_ERR_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("TilingForMoeGatingTopKHash", "Tiling context is null"),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
auto platformInfo = context->GetPlatformInfo();
|
||||
OPS_ERR_IF(platformInfo == nullptr, OPS_REPORT_VECTOR_INNER_ERR("TilingForMoeGatingTopKHash", "Tiling platformInfo is null"),
|
||||
return ge::GRAPH_FAILED);
|
||||
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
|
||||
auto socVersion = ascendcPlatform.GetSocVersion();
|
||||
if (socVersion == platform_ascendc::SocVersion::ASCEND950) {
|
||||
OPS_LOG_I(context, "Using arch35 tiling for ASCEND950");
|
||||
MoeGatingTopKHashRegBase::MoeGatingTopKHashTilingRegbase moeGatingTopKTilingRegbase(context);
|
||||
return moeGatingTopKTilingRegbase.DoOpTiling();
|
||||
}
|
||||
|
||||
MoeGatingTopKHashTilingBase moeGatingTopKTilingBase(context);
|
||||
return moeGatingTopKTilingBase.DoOpTiling();
|
||||
}
|
||||
|
||||
static ge::graphStatus TilingPrepareForMoeGatingTopKHash(gert::TilingParseContext *context)
|
||||
{
|
||||
(void)context;
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
IMPL_OP_OPTILING(MoeGatingTopKHash)
|
||||
.Tiling(TilingForMoeGatingTopKHash)
|
||||
.TilingParse<MoeGatingTopKHashCompileInfo>(TilingPrepareForMoeGatingTopKHash);
|
||||
} // namespace optiling
|
||||
@@ -0,0 +1,84 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file moe_gating_top_k_hash_tiling.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_MOE_GATING_TOP_K_H
|
||||
#define AIR_CXX_RUNTIME_V2_OP_IMPL_MOE_GATING_TOP_K_H
|
||||
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <vector>
|
||||
#include <algorithm>
|
||||
#include <iostream>
|
||||
#include "register/op_impl_registry.h"
|
||||
#include "platform/platform_infos_def.h"
|
||||
#include "exe_graph/runtime/tiling_context.h"
|
||||
#include "tiling/platform/platform_ascendc.h"
|
||||
#include "register/op_def_registry.h"
|
||||
#include "register/tilingdata_base.h"
|
||||
#include "tiling/tiling_api.h"
|
||||
#include "error/ops_error.h"
|
||||
#include "platform/platform_info.h"
|
||||
|
||||
namespace optiling {
|
||||
BEGIN_TILING_DATA_DEF(MoeGatingTopKHashTilingData)
|
||||
TILING_DATA_FIELD_DEF(int64_t, needCoreNum);
|
||||
TILING_DATA_FIELD_DEF(int64_t, rowCount);
|
||||
TILING_DATA_FIELD_DEF(int64_t, perCoreRowCount);
|
||||
TILING_DATA_FIELD_DEF(int64_t, lastCoreRowCount);
|
||||
TILING_DATA_FIELD_DEF(int64_t, expertCount);
|
||||
TILING_DATA_FIELD_DEF(int64_t, addBias);
|
||||
TILING_DATA_FIELD_DEF(int64_t, k);
|
||||
TILING_DATA_FIELD_DEF(int64_t, kGroup);
|
||||
TILING_DATA_FIELD_DEF(int64_t, groupCount);
|
||||
TILING_DATA_FIELD_DEF(int64_t, perGroupExpertCount);
|
||||
TILING_DATA_FIELD_DEF(int64_t, perGroupExpertCountAlign);
|
||||
TILING_DATA_FIELD_DEF(int64_t, groupSelectMode);
|
||||
TILING_DATA_FIELD_DEF(int64_t, renorm);
|
||||
TILING_DATA_FIELD_DEF(int64_t, normType);
|
||||
TILING_DATA_FIELD_DEF(int64_t, outFlag);
|
||||
TILING_DATA_FIELD_DEF(int64_t, hashFlag);
|
||||
TILING_DATA_FIELD_DEF(int64_t, vmsCount);
|
||||
TILING_DATA_FIELD_DEF(float, routedScalingFactor);
|
||||
TILING_DATA_FIELD_DEF(float, eps);
|
||||
TILING_DATA_FIELD_DEF(int64_t, calTmpBufUbSize);
|
||||
END_TILING_DATA_DEF;
|
||||
REGISTER_TILING_DATA_CLASS(MoeGatingTopKHash, MoeGatingTopKHashTilingData)
|
||||
|
||||
BEGIN_TILING_DATA_DEF(MoeGatingTopKHashRegbaseTilingData)
|
||||
TILING_DATA_FIELD_DEF(int64_t, needCoreNum);
|
||||
TILING_DATA_FIELD_DEF(int64_t, rowCount);
|
||||
TILING_DATA_FIELD_DEF(int64_t, perCoreRowCount);
|
||||
TILING_DATA_FIELD_DEF(int64_t, lastCoreRowCount);
|
||||
TILING_DATA_FIELD_DEF(int64_t, expertCount);
|
||||
TILING_DATA_FIELD_DEF(int64_t, addBias);
|
||||
TILING_DATA_FIELD_DEF(int64_t, k);
|
||||
TILING_DATA_FIELD_DEF(int64_t, kGroup);
|
||||
TILING_DATA_FIELD_DEF(int64_t, groupCount);
|
||||
TILING_DATA_FIELD_DEF(int64_t, perGroupExpertCount);
|
||||
TILING_DATA_FIELD_DEF(int64_t, perGroupExpertCountAlign);
|
||||
TILING_DATA_FIELD_DEF(int64_t, groupSelectMode);
|
||||
TILING_DATA_FIELD_DEF(int64_t, renorm);
|
||||
TILING_DATA_FIELD_DEF(int64_t, normType);
|
||||
TILING_DATA_FIELD_DEF(int64_t, outFlag);
|
||||
TILING_DATA_FIELD_DEF(int64_t, hashFlag);
|
||||
TILING_DATA_FIELD_DEF(int64_t, vmsCount);
|
||||
TILING_DATA_FIELD_DEF(float, routedScalingFactor);
|
||||
TILING_DATA_FIELD_DEF(float, eps);
|
||||
TILING_DATA_FIELD_DEF_STRUCT(SoftMaxTiling, softmaxTilingData);
|
||||
END_TILING_DATA_DEF;
|
||||
REGISTER_TILING_DATA_CLASS(MoeGatingTopKHash_10000, MoeGatingTopKHashRegbaseTilingData)
|
||||
struct MoeGatingTopKHashCompileInfo {};
|
||||
} // namespace optiling
|
||||
#endif // AIR_CXX_RUNTIME_V2_OP_IMPL_MOE_GATING_TOP_K_H
|
||||
@@ -0,0 +1,594 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/* !
|
||||
* \file moe_gating_top_k_hash_tiling_arch35.cpp
|
||||
* \brief
|
||||
*/
|
||||
#include "moe_gating_top_k_hash_tiling.h"
|
||||
#include <graph/utils/type_utils.h>
|
||||
|
||||
namespace optiling {
|
||||
namespace MoeGatingTopKHashRegBase {
|
||||
|
||||
const static uint64_t MOE_GATING_TOP_K_REGBASE_TILING_KEY = 10000;
|
||||
const static uint64_t MOE_GATING_TOP_K_REGBASE_TILING_KEY_1 = 10001;
|
||||
const static uint64_t MOE_GATING_TOP_K_REGBASE_TILING_KEY_2 = 10002;
|
||||
const static uint64_t MOE_GATING_TOP_K_REGBASE_TILING_KEY_3 = 10003;
|
||||
const static uint64_t MOE_GATING_TOP_K_REGBASE_TILING_KEY_4 = 10004;
|
||||
|
||||
const static int64_t GROUP_SELECT_MODE_MAX = 0;
|
||||
const static int64_t GROUP_SELECT_MODE_SUM = 1;
|
||||
const static int64_t RENORM_NO = 0;
|
||||
const static int64_t RENORM_L1 = 1;
|
||||
const static int64_t NORM_TYPE_SOFTMAX = 0;
|
||||
const static int64_t NORM_TYPE_SIGMOID = 1;
|
||||
const static int64_t NORM_TYPE_SOFTPLUS = 2;
|
||||
const static int64_t OUT_FLAG_FALSE = 0;
|
||||
const static int64_t OUT_FLAG_TRUE = 1;
|
||||
const static size_t X_INPUT_DIMS = 2;
|
||||
const static size_t BIAS_INPUT_DIMS = 1;
|
||||
const static size_t Y_OUTPUT_DIMS = 2;
|
||||
const static size_t EXPERT_IDX_OUTPUY_DIMS = 2;
|
||||
const static size_t OUT_OUTPUT_DIMS = 2;
|
||||
const static int64_t MAX_EXPERT_COUNT = 2048;
|
||||
|
||||
const static int64_t X_INPUT_INDEX = 0;
|
||||
const static int64_t BIAS_INPUT_INDEX = 1;
|
||||
const static int64_t INPUT_IDS_INPUT_INDEX = 2;
|
||||
const static int64_t TID_TO_EID_INPUT_INDEX = 3;
|
||||
const static int64_t Y_OUTPUT_INDEX = 0;
|
||||
const static int64_t EXPERT_IDX_OUTPUT_INDEX = 1;
|
||||
const static int64_t OUT_OUTPUT_INDEX = 2;
|
||||
const static int64_t K_ATTR_INDEX = 0;
|
||||
const static int64_t K_GROUP_ATTR_INDEX = 1;
|
||||
const static int64_t GROUP_COUNT_ATTR_INDEX = 2;
|
||||
const static int64_t GROUP_SELECT_MODE_ATTR_INDEX = 3;
|
||||
const static int64_t RENORM_ATTR_INDEX = 4;
|
||||
const static int64_t MRGSORT_SIZE = 4;
|
||||
const static int64_t NORM_TYPE_ATTR_INDEX = 5;
|
||||
const static int64_t OUT_FLAG_ATTR_INDEX = 6;
|
||||
const static int64_t ROUTED_SCALING_FACTOR_ATTR_INDEX = 7;
|
||||
const static int64_t EPS_ATTR_INDEX = 8;
|
||||
const static int64_t DEFAULT_WORKSPACE_SIZE = static_cast<int64_t>(16 * 1024 * 1024); // 预留16M空间
|
||||
|
||||
template <typename T>
|
||||
static inline T CeilDiv(T num, T rnd)
|
||||
{
|
||||
return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static inline T CeilAlign(T num, T rnd)
|
||||
{
|
||||
return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)) * (rnd));
|
||||
}
|
||||
|
||||
class MoeGatingTopKHashTilingRegbase {
|
||||
public:
|
||||
explicit MoeGatingTopKHashTilingRegbase(gert::TilingContext *context) : context_(context)
|
||||
{
|
||||
Reset();
|
||||
}
|
||||
~MoeGatingTopKHashTilingRegbase() = default;
|
||||
|
||||
bool IsCapable()
|
||||
{
|
||||
if (socVersion_ != platform_ascendc::SocVersion::ASCEND950) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
// 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小
|
||||
ge::graphStatus GetPlatformInfo();
|
||||
// 2、获取INPUT/OUTPUT/ATTR信息
|
||||
ge::graphStatus GetShapeAttrsInfo();
|
||||
// 3、计算数据切分TilingData
|
||||
ge::graphStatus DoOpTiling();
|
||||
// 4、计算高阶API的TilingData
|
||||
ge::graphStatus DoLibApiTiling();
|
||||
// 5、计算TilingKey
|
||||
uint64_t GetTilingKey() const;
|
||||
// 6、计算Workspace 大小
|
||||
ge::graphStatus GetWorkspaceSize();
|
||||
// 7、保存Tiling数据
|
||||
ge::graphStatus PostTiling();
|
||||
void Reset();
|
||||
|
||||
private:
|
||||
ge::graphStatus CheckInputShape();
|
||||
ge::graphStatus CheckAttr();
|
||||
ge::graphStatus CheckOutShape();
|
||||
void CalTmpBufUbSize();
|
||||
void SplitRows();
|
||||
void Tiling4GatherOutComputeSplitK();
|
||||
|
||||
const gert::Shape *xShape_ = nullptr;
|
||||
const gert::Shape *biasShape_ = nullptr;
|
||||
const gert::Shape *yShape_ = nullptr;
|
||||
const gert::Shape *expertIdxShape_ = nullptr;
|
||||
const gert::Shape *outShape_ = nullptr;
|
||||
const gert::Shape *inputIdsShape_ = nullptr;
|
||||
const gert::Shape *tid2eidShape_ = nullptr;
|
||||
|
||||
ge::DataType inputIdsDtype;
|
||||
ge::DataType tid2eidDtype;
|
||||
|
||||
uint64_t coreNum_ = 0;
|
||||
int64_t rows_ = 0;
|
||||
int64_t expertCount_ = 0;
|
||||
int64_t addBias_ = 0;
|
||||
|
||||
int64_t k_ = 0;
|
||||
int64_t kGroup_ = 1;
|
||||
int64_t groupCount_ = 1;
|
||||
int64_t groupSelectMode_ = GROUP_SELECT_MODE_MAX;
|
||||
int64_t renorm_ = RENORM_NO;
|
||||
int64_t normType_ = NORM_TYPE_SOFTMAX;
|
||||
int64_t outFlag_ = OUT_FLAG_FALSE;
|
||||
int64_t hashFlag_ = 0;
|
||||
float routedScalingFactor_ = 1.0;
|
||||
float eps_ = 1e-20f;
|
||||
|
||||
int64_t inputDtypeSize_;
|
||||
const char *opName_ = "";
|
||||
MoeGatingTopKHashRegbaseTilingData moeGatingTopKTilingData_;
|
||||
gert::TilingContext *context_ = nullptr;
|
||||
uint64_t workspaceSize_ = 0;
|
||||
platform_ascendc::SocVersion socVersion_;
|
||||
};
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingRegbase::CheckInputShape()
|
||||
{
|
||||
size_t xDimNum = xShape_->GetDimNum();
|
||||
OPS_ERR_IF(xDimNum != X_INPUT_DIMS,
|
||||
OPS_LOG_E(context_, "The dim number of x is: %zu, but should be %zu.", xDimNum, X_INPUT_DIMS),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
// 通过输入获取rows 和 expertCount
|
||||
rows_ = xShape_->GetDim(0);
|
||||
expertCount_ = xShape_->GetDim(1);
|
||||
moeGatingTopKTilingData_.set_rowCount(rows_);
|
||||
moeGatingTopKTilingData_.set_expertCount(expertCount_);
|
||||
OPS_ERR_IF(
|
||||
expertCount_ > MAX_EXPERT_COUNT,
|
||||
OPS_LOG_E(context_, "expert count is: %ld, but should not greater than %ld.", expertCount_, MAX_EXPERT_COUNT),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
if (biasShape_ != nullptr) {
|
||||
addBias_ = 1;
|
||||
size_t biasDimNum = biasShape_->GetDimNum();
|
||||
OPS_ERR_IF(biasDimNum != BIAS_INPUT_DIMS,
|
||||
OPS_LOG_E(context_, "The number of bias dim is: %zu, but should be %zu.", biasDimNum, BIAS_INPUT_DIMS),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF(biasShape_->GetDim(0) != expertCount_,
|
||||
OPS_LOG_E(context_, "The first dim of bias is: %ld, but should be expert num: %ld.",
|
||||
biasShape_->GetDim(0), expertCount_),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
moeGatingTopKTilingData_.set_addBias(addBias_);
|
||||
|
||||
if (inputIdsShape_ != nullptr) {
|
||||
OPS_ERR_IF(
|
||||
tid2eidShape_ == nullptr,
|
||||
OPS_LOG_E(context_, "The tid2eid should not be empty when inputIds has value."),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
if (tid2eidShape_ != nullptr) {
|
||||
OPS_ERR_IF(
|
||||
inputIdsShape_ == nullptr,
|
||||
OPS_LOG_E(context_, "The inputIds should not be empty when tid2eid has value."),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
|
||||
if (inputIdsShape_ != nullptr && tid2eidShape_ != nullptr) {
|
||||
hashFlag_ = 1;
|
||||
OPS_LOG_I(context_, "hashFlag_ is 1.");
|
||||
}
|
||||
moeGatingTopKTilingData_.set_hashFlag(hashFlag_);
|
||||
OPS_ERR_IF(k_ > expertCount_,
|
||||
OPS_LOG_E(context_, "k is: %ld, expert num is: %ld, k cannot be greater than expert num.", k_, expertCount_),
|
||||
return ge::GRAPH_FAILED);
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingRegbase::CheckAttr()
|
||||
{
|
||||
OPS_ERR_IF(k_ <= 0, OPS_LOG_E(context_, "k is: %ld, but should be greater than 0.", k_), return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF(kGroup_ <= 0, OPS_LOG_E(context_, "k_group is: %ld, but should be greater than 0.", kGroup_),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF(groupCount_ <= 0, OPS_LOG_E(context_, "group_count is: %ld, but should be greater than 0.", groupCount_),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF(expertCount_ % groupCount_ != 0,
|
||||
OPS_LOG_E(context_, "expert num : %ld is not divisible by group_count: %ld", expertCount_, groupCount_),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF(kGroup_ > groupCount_,
|
||||
OPS_LOG_E(context_, "k_group is: %ld, but should not greater than group_count: %ld", kGroup_, groupCount_),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF(groupCount_ == expertCount_ && kGroup_ < k_,
|
||||
OPS_LOG_E(context_, "k_group * group expert count is: %ld, but it must be greater than or equal to k: %ld.",
|
||||
kGroup_, k_),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
if (kGroup_ == groupCount_ || groupCount_ == expertCount_) {
|
||||
kGroup_ = 1;
|
||||
groupCount_ = 1;
|
||||
}
|
||||
moeGatingTopKTilingData_.set_kGroup(kGroup_);
|
||||
moeGatingTopKTilingData_.set_groupCount(groupCount_);
|
||||
int64_t groupExpertCount = expertCount_ / groupCount_;
|
||||
int64_t groupExpertCountAlign = CeilAlign(groupExpertCount, 32L);
|
||||
moeGatingTopKTilingData_.set_perGroupExpertCount(expertCount_ / groupCount_);
|
||||
moeGatingTopKTilingData_.set_perGroupExpertCountAlign(groupExpertCountAlign);
|
||||
|
||||
OPS_ERR_IF(groupCount_ * groupExpertCountAlign > MAX_EXPERT_COUNT,
|
||||
OPS_LOG_E(context_, "group count * group expert count align is: %ld, but should not greater than %ld.",
|
||||
groupCount_ * groupExpertCountAlign, MAX_EXPERT_COUNT),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
OPS_ERR_IF(kGroup_ * groupExpertCount < k_,
|
||||
OPS_LOG_E(context_, "k_group * group expert count is: %ld, but it must be greater than or equal to k: %ld.",
|
||||
kGroup_ * groupExpertCount, k_),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
OPS_ERR_IF(groupExpertCount < 1,
|
||||
OPS_LOG_E(context_, "per group expert count is: %ld, but should be greater than 0.", groupExpertCount),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF(
|
||||
groupSelectMode_ != GROUP_SELECT_MODE_SUM && groupSelectMode_ != GROUP_SELECT_MODE_MAX,
|
||||
OPS_LOG_E(context_, "group select mode is: %ld, but currently only support %ld and %ld.", groupSelectMode_,
|
||||
GROUP_SELECT_MODE_SUM, GROUP_SELECT_MODE_MAX),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF(groupSelectMode_ == GROUP_SELECT_MODE_SUM && groupExpertCount < 2,
|
||||
OPS_LOG_E(context_,
|
||||
"group expert count is: %ld, if group select mode is: %ld, group expert count should be greater than 1.",
|
||||
groupExpertCount, groupSelectMode_),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
OPS_ERR_IF(renorm_ != RENORM_NO,
|
||||
OPS_LOG_E(context_, "renorm is: %ld, but currently only support %ld.", renorm_, RENORM_NO),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
OPS_ERR_IF(normType_ != NORM_TYPE_SOFTMAX && normType_ != NORM_TYPE_SIGMOID && normType_ != NORM_TYPE_SOFTPLUS,
|
||||
OPS_LOG_E(context_, "norm type is: %ld, but currently only support %ld, %ld and %ld.", normType_,
|
||||
NORM_TYPE_SOFTMAX, NORM_TYPE_SIGMOID, NORM_TYPE_SOFTPLUS),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingRegbase::GetShapeAttrsInfo()
|
||||
{
|
||||
opName_ = context_->GetNodeName();
|
||||
// 获取输入shape信息
|
||||
auto xShapePtr = context_->GetInputShape(X_INPUT_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, xShapePtr, return ge::GRAPH_FAILED);
|
||||
xShape_ = &xShapePtr->GetStorageShape();
|
||||
auto biasShapePtr = context_->GetOptionalInputShape(BIAS_INPUT_INDEX);
|
||||
biasShape_ = biasShapePtr == nullptr ? nullptr : &biasShapePtr->GetStorageShape();
|
||||
auto inputIdsShapePtr = context_->GetOptionalInputShape(INPUT_IDS_INPUT_INDEX);
|
||||
inputIdsShape_ = inputIdsShapePtr == nullptr ? nullptr : &inputIdsShapePtr->GetStorageShape();
|
||||
auto tid2eidShapePtr = context_->GetOptionalInputShape(TID_TO_EID_INPUT_INDEX);
|
||||
tid2eidShape_ = tid2eidShapePtr == nullptr ? nullptr : &tid2eidShapePtr->GetStorageShape();
|
||||
|
||||
// 获取输出shape
|
||||
auto yShapePtr = context_->GetOutputShape(Y_OUTPUT_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, yShapePtr, return ge::GRAPH_FAILED);
|
||||
yShape_ = &yShapePtr->GetStorageShape();
|
||||
auto expertIdxPtr = context_->GetOutputShape(EXPERT_IDX_OUTPUT_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, expertIdxPtr, return ge::GRAPH_FAILED);
|
||||
expertIdxShape_ = &expertIdxPtr->GetStorageShape();
|
||||
auto outPtr = context_->GetOutputShape(OUT_OUTPUT_INDEX);
|
||||
if (outPtr != nullptr) {
|
||||
outShape_ = &outPtr->GetStorageShape();
|
||||
}
|
||||
|
||||
auto x = context_->GetInputDesc(X_INPUT_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, x, return ge::GRAPH_FAILED);
|
||||
auto xDtype = x->GetDataType();
|
||||
OPS_ERR_IF(
|
||||
(xDtype != ge::DataType::DT_FLOAT && xDtype != ge::DataType::DT_FLOAT16 && xDtype != ge::DataType::DT_BF16),
|
||||
OPS_LOG_E(context_, "x dtype %s error, only supports float32, half, bf16. please check.",
|
||||
ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
if (biasShapePtr != nullptr) {
|
||||
auto biasDtype = context_->GetOptionalInputDesc(BIAS_INPUT_INDEX)->GetDataType();
|
||||
OPS_ERR_IF((biasDtype != xDtype),
|
||||
OPS_LOG_E(context_, "bias dtype %s not equal x dtype %s, please check.",
|
||||
ge::TypeUtils::DataTypeToSerialString(biasDtype).c_str(),
|
||||
ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
if (inputIdsShapePtr != nullptr) {
|
||||
inputIdsDtype = context_->GetOptionalInputDesc(INPUT_IDS_INPUT_INDEX)->GetDataType();
|
||||
OPS_ERR_IF((inputIdsDtype != ge::DataType::DT_INT32 && inputIdsDtype != ge::DataType::DT_INT64),
|
||||
OPS_LOG_E(context_, "inputIds dtype %s error, only supports int32 and int64. please check.",
|
||||
ge::TypeUtils::DataTypeToSerialString(inputIdsDtype).c_str()),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
if (tid2eidShapePtr != nullptr) {
|
||||
tid2eidDtype = context_->GetOptionalInputDesc(TID_TO_EID_INPUT_INDEX)->GetDataType();
|
||||
OPS_ERR_IF((tid2eidDtype != ge::DataType::DT_INT32 && tid2eidDtype != ge::DataType::DT_INT64),
|
||||
OPS_LOG_E(context_, "tid2eid dtype %s error, only supports int32 and int64. please check.",
|
||||
ge::TypeUtils::DataTypeToSerialString(tid2eidDtype).c_str()),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
|
||||
auto yDesc = context_->GetOutputDesc(Y_OUTPUT_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, yDesc, return ge::GRAPH_FAILED);
|
||||
auto yDtype = yDesc->GetDataType();
|
||||
OPS_ERR_IF((yDtype != xDtype),
|
||||
OPS_LOG_E(context_, "y out dtype %s must be the same with x dtype %s.",
|
||||
ge::TypeUtils::DataTypeToSerialString(yDtype).c_str(),
|
||||
ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
auto expertIdDesc = context_->GetOutputDesc(EXPERT_IDX_OUTPUT_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, expertIdDesc, return ge::GRAPH_FAILED);
|
||||
auto expertIdDtype = expertIdDesc->GetDataType();
|
||||
OPS_ERR_IF((expertIdDtype != ge::DataType::DT_INT32),
|
||||
OPS_LOG_E(context_, "expertId out dtype %s error, only supports int32. please check.",
|
||||
ge::TypeUtils::DataTypeToSerialString(expertIdDtype).c_str()),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
// 获取属性
|
||||
auto attrs = context_->GetAttrs();
|
||||
OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED);
|
||||
|
||||
const int64_t *kPtr = attrs->GetAttrPointer<int64_t>(K_ATTR_INDEX);
|
||||
OPS_LOG_E_IF_NULL(context_, kPtr, return ge::GRAPH_FAILED);
|
||||
k_ = *kPtr;
|
||||
moeGatingTopKTilingData_.set_k(k_);
|
||||
OPS_LOG_I(context_, "Attr k is: %ld ", k_);
|
||||
|
||||
const int64_t *kGroupPtr = attrs->GetAttrPointer<int64_t>(K_GROUP_ATTR_INDEX);
|
||||
if (kGroupPtr != nullptr) {
|
||||
kGroup_ = *kGroupPtr;
|
||||
}
|
||||
OPS_LOG_I(context_, "Attr k_group is: %ld ", kGroup_);
|
||||
|
||||
const int64_t *groupCountPtr = attrs->GetAttrPointer<int64_t>(GROUP_COUNT_ATTR_INDEX);
|
||||
if (groupCountPtr != nullptr) {
|
||||
groupCount_ = *groupCountPtr;
|
||||
}
|
||||
OPS_LOG_I(context_, "Attr group_count is: %ld ", groupCount_);
|
||||
|
||||
const int64_t *groupSelectModePtr = attrs->GetAttrPointer<int64_t>(GROUP_SELECT_MODE_ATTR_INDEX);
|
||||
if (groupSelectModePtr != nullptr) {
|
||||
groupSelectMode_ = *groupSelectModePtr;
|
||||
}
|
||||
moeGatingTopKTilingData_.set_groupSelectMode(groupSelectMode_);
|
||||
OPS_LOG_I(context_, "Attr group_select_mode is: %ld ", groupSelectMode_);
|
||||
|
||||
const int64_t *renormPtr = attrs->GetAttrPointer<int64_t>(RENORM_ATTR_INDEX);
|
||||
if (renormPtr != nullptr) {
|
||||
renorm_ = *renormPtr;
|
||||
}
|
||||
moeGatingTopKTilingData_.set_renorm(renorm_);
|
||||
OPS_LOG_I(context_, "Attr renorm is: %ld ", renorm_);
|
||||
|
||||
const int64_t *normTypePtr = attrs->GetAttrPointer<int64_t>(NORM_TYPE_ATTR_INDEX);
|
||||
if (normTypePtr != nullptr) {
|
||||
normType_ = *normTypePtr;
|
||||
}
|
||||
moeGatingTopKTilingData_.set_normType(normType_);
|
||||
OPS_LOG_I(context_, "Attr norm_type is: %ld ", normType_);
|
||||
|
||||
const bool *outFlagPtr = attrs->GetAttrPointer<bool>(OUT_FLAG_ATTR_INDEX);
|
||||
if (outFlagPtr != nullptr) {
|
||||
outFlag_ = (*outFlagPtr) ? 1 : 0;
|
||||
}
|
||||
moeGatingTopKTilingData_.set_outFlag(outFlag_);
|
||||
OPS_LOG_I(context_, "Attr out_flag is: %ld ", outFlag_);
|
||||
|
||||
const float *routedScalingFactorPtr = attrs->GetAttrPointer<float>(ROUTED_SCALING_FACTOR_ATTR_INDEX);
|
||||
if (routedScalingFactorPtr != nullptr) {
|
||||
routedScalingFactor_ = *routedScalingFactorPtr;
|
||||
}
|
||||
moeGatingTopKTilingData_.set_routedScalingFactor(routedScalingFactor_);
|
||||
OPS_LOG_I(context_, "Attr routed_scaling_factor is: %f ", routedScalingFactor_);
|
||||
|
||||
const float *epsPtr = attrs->GetAttrPointer<float>(EPS_ATTR_INDEX);
|
||||
if (epsPtr != nullptr) {
|
||||
eps_ = *epsPtr;
|
||||
}
|
||||
moeGatingTopKTilingData_.set_eps(eps_);
|
||||
OPS_LOG_I(context_, "Attr eps is: %f ", eps_);
|
||||
|
||||
auto outDesc = context_->GetOutputDesc(OUT_OUTPUT_INDEX);
|
||||
if (outFlag_ && outDesc != nullptr) {
|
||||
auto outDtype = outDesc->GetDataType();
|
||||
OPS_ERR_IF((outDtype != ge::DataType::DT_FLOAT),
|
||||
OPS_LOG_E(context_, "norm out dtype %s error, only supports float32. please check.",
|
||||
ge::TypeUtils::DataTypeToSerialString(outDtype).c_str()),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
|
||||
inputDtypeSize_ = static_cast<int64_t>(ge::GetSizeByDataType(context_->GetInputDesc(0)->GetDataType()));
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingRegbase::GetPlatformInfo()
|
||||
{
|
||||
auto platformInfo = context_->GetPlatformInfo();
|
||||
OPS_ERR_IF(platformInfo == nullptr, OPS_LOG_E(context_, "fail to get platform info"), return ge::GRAPH_FAILED);
|
||||
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
|
||||
coreNum_ = ascendcPlatform.GetCoreNumAiv();
|
||||
socVersion_ = ascendcPlatform.GetSocVersion();
|
||||
uint64_t ubSizePlatForm;
|
||||
ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm);
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingRegbase::CheckOutShape()
|
||||
{
|
||||
OPS_ERR_IF((yShape_->GetDimNum() != xShape_->GetDimNum()),
|
||||
OPS_LOG_E(context_, "y out shape num %zu and x shape num %zu not equal, please check.", yShape_->GetDimNum(),
|
||||
xShape_->GetDimNum()),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF((expertIdxShape_->GetDimNum() != xShape_->GetDimNum()),
|
||||
OPS_LOG_E(context_, "expertId out shape num %zu and x shape num %zu not equal, please check.",
|
||||
expertIdxShape_->GetDimNum(), xShape_->GetDimNum()),
|
||||
return ge::GRAPH_FAILED);
|
||||
if (outShape_ != nullptr) {
|
||||
OPS_ERR_IF((outShape_->GetDimNum() != xShape_->GetDimNum()),
|
||||
OPS_LOG_E(context_, "norm out shape num %zu and x shape num %zu not equal, please check.",
|
||||
outShape_->GetDimNum(), xShape_->GetDimNum()),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
|
||||
OPS_ERR_IF((yShape_->GetDim(0) != xShape_->GetDim(0)),
|
||||
OPS_LOG_E(context_, "y out dim[0] %ld not equal x dim[0] %ld, please check.", yShape_->GetDim(0),
|
||||
xShape_->GetDim(0)),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF((expertIdxShape_->GetDim(0) != xShape_->GetDim(0)),
|
||||
OPS_LOG_E(context_, "expertId out dim[0] %ld not equal x dim[0] %ld, please check.",
|
||||
expertIdxShape_->GetDim(0), xShape_->GetDim(0)),
|
||||
return ge::GRAPH_FAILED);
|
||||
if (outFlag_ && outShape_ != nullptr) {
|
||||
OPS_ERR_IF((outShape_->GetDim(0) != xShape_->GetDim(0)),
|
||||
OPS_LOG_E(context_, "norm out dim[0] %ld and x dim[0] %ld not equal, please check.",
|
||||
outShape_->GetDim(0), outShape_->GetDim(0)),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
|
||||
OPS_ERR_IF((yShape_->GetDim(1) != k_),
|
||||
OPS_LOG_E(context_, "y dim[1] %ld not equal k %ld, please check.", yShape_->GetDim(1), k_),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_ERR_IF((expertIdxShape_->GetDim(1) != k_),
|
||||
OPS_LOG_E(context_, "expertId dim[1] %ld not equal k %ld, please check.", expertIdxShape_->GetDim(1), k_),
|
||||
return ge::GRAPH_FAILED);
|
||||
if (outFlag_ && outShape_ != nullptr) {
|
||||
OPS_ERR_IF((outShape_->GetDim(1) != xShape_->GetDim(1)),
|
||||
OPS_LOG_E(context_, "normOut dim[1] %ld and x dim[1] %ld not equal, please check.", outShape_->GetDim(1),
|
||||
xShape_->GetDim(1)),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
void MoeGatingTopKHashTilingRegbase::CalTmpBufUbSize()
|
||||
{
|
||||
std::vector<int64_t> shape_vec = {groupCount_ * moeGatingTopKTilingData_.get_perGroupExpertCountAlign()};
|
||||
ge::Shape softmaxShape(shape_vec);
|
||||
|
||||
uint32_t softmaxTmpSize = AscendC::GetSoftMaxMaxTmpSize(softmaxShape, sizeof(float), true);
|
||||
AscendC::SoftMaxTilingFunc(softmaxShape, sizeof(float), softmaxTmpSize, moeGatingTopKTilingData_.softmaxTilingData);
|
||||
}
|
||||
|
||||
void MoeGatingTopKHashTilingRegbase::SplitRows()
|
||||
{
|
||||
int64_t perCoreRows = CeilDiv(rows_, static_cast<int64_t>(coreNum_));
|
||||
int64_t needCoreNum = CeilDiv(rows_, perCoreRows);
|
||||
if (perCoreRows == 0) {
|
||||
OPS_LOG_E(context_, "perCoreRows can't be 0.");
|
||||
return;
|
||||
}
|
||||
int64_t lastCoreRows = rows_ % perCoreRows == 0 ? perCoreRows : rows_ % perCoreRows;
|
||||
moeGatingTopKTilingData_.set_needCoreNum(needCoreNum);
|
||||
moeGatingTopKTilingData_.set_perCoreRowCount(perCoreRows);
|
||||
moeGatingTopKTilingData_.set_lastCoreRowCount(lastCoreRows);
|
||||
|
||||
int64_t vmsCount = 0;
|
||||
if (kGroup_ > MRGSORT_SIZE) {
|
||||
int64_t index = MRGSORT_SIZE;
|
||||
while (index < kGroup_) {
|
||||
index = index * MRGSORT_SIZE;
|
||||
vmsCount++;
|
||||
}
|
||||
}
|
||||
moeGatingTopKTilingData_.set_vmsCount(vmsCount);
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingRegbase::DoOpTiling()
|
||||
{
|
||||
auto ret = GetPlatformInfo();
|
||||
if (ret != ge::GRAPH_SUCCESS) {
|
||||
return ret;
|
||||
}
|
||||
|
||||
ret = GetShapeAttrsInfo();
|
||||
if (ret != ge::GRAPH_SUCCESS) {
|
||||
return ret;
|
||||
}
|
||||
|
||||
ret = CheckInputShape();
|
||||
if (ret != ge::GRAPH_SUCCESS) {
|
||||
return ret;
|
||||
}
|
||||
|
||||
ret = CheckAttr();
|
||||
if (ret != ge::GRAPH_SUCCESS) {
|
||||
return ret;
|
||||
}
|
||||
|
||||
ret = CheckOutShape();
|
||||
if (ret != ge::GRAPH_SUCCESS) {
|
||||
return ret;
|
||||
}
|
||||
|
||||
CalTmpBufUbSize();
|
||||
SplitRows();
|
||||
|
||||
ret = PostTiling();
|
||||
if (ret != ge::GRAPH_SUCCESS) {
|
||||
return ret;
|
||||
}
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingRegbase::DoLibApiTiling()
|
||||
{
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingRegbase::GetWorkspaceSize()
|
||||
{
|
||||
// 计算workspace大小
|
||||
workspaceSize_ = DEFAULT_WORKSPACE_SIZE;
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus MoeGatingTopKHashTilingRegbase::PostTiling()
|
||||
{
|
||||
context_->SetTilingKey(GetTilingKey());
|
||||
context_->SetBlockDim(moeGatingTopKTilingData_.get_needCoreNum());
|
||||
size_t *currentWorkspace = context_->GetWorkspaceSizes(1);
|
||||
currentWorkspace[0] = workspaceSize_;
|
||||
moeGatingTopKTilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(),
|
||||
context_->GetRawTilingData()->GetCapacity());
|
||||
context_->GetRawTilingData()->SetDataSize(moeGatingTopKTilingData_.GetDataSize());
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
uint64_t MoeGatingTopKHashTilingRegbase::GetTilingKey() const
|
||||
{
|
||||
if (inputIdsShape_ == nullptr) {
|
||||
return MOE_GATING_TOP_K_REGBASE_TILING_KEY;
|
||||
} else if (inputIdsDtype == ge::DataType::DT_INT32 && tid2eidDtype == ge::DataType::DT_INT64) {
|
||||
return MOE_GATING_TOP_K_REGBASE_TILING_KEY_1;
|
||||
} else if (inputIdsDtype == ge::DataType::DT_INT32 && tid2eidDtype == ge::DataType::DT_INT32) {
|
||||
return MOE_GATING_TOP_K_REGBASE_TILING_KEY_2;
|
||||
} else if (inputIdsDtype == ge::DataType::DT_INT64 && tid2eidDtype == ge::DataType::DT_INT64) {
|
||||
return MOE_GATING_TOP_K_REGBASE_TILING_KEY_3;
|
||||
} else if (inputIdsDtype == ge::DataType::DT_INT64 && tid2eidDtype == ge::DataType::DT_INT32) {
|
||||
return MOE_GATING_TOP_K_REGBASE_TILING_KEY_4;
|
||||
}
|
||||
}
|
||||
|
||||
void MoeGatingTopKHashTilingRegbase::Reset()
|
||||
{
|
||||
opName_ = nullptr;
|
||||
return;
|
||||
}
|
||||
|
||||
} // namespace MoeGatingTopKHashRegBase
|
||||
} // namespace optiling
|
||||
214
csrc/moe/moe_gating_top_k_hash/op_kernel/arch35/common.h
Normal file
214
csrc/moe/moe_gating_top_k_hash/op_kernel/arch35/common.h
Normal file
@@ -0,0 +1,214 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file common.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef MOE_GATING_TOP_K_COMMON_H_V35
|
||||
#define MOE_GATING_TOP_K_COMMON_H_V35
|
||||
|
||||
#include "kernel_operator.h"
|
||||
|
||||
namespace MoeGatingTopKHashRegbaseNS {
|
||||
using namespace AscendC;
|
||||
constexpr int32_t MIN_FP32 = 0xFF800000;
|
||||
constexpr int64_t ONE_REPEAT_SORT_NUM = 32;
|
||||
constexpr int64_t BLOCK_BYTES = 32;
|
||||
|
||||
constexpr int64_t MERGE_LIST_TWO = 2;
|
||||
constexpr int64_t MERGE_LIST_THREE = 3;
|
||||
constexpr int64_t MERGE_LIST_FOUR = 4;
|
||||
|
||||
constexpr int64_t MERGE_LIST_IDX_TWO = 2;
|
||||
constexpr int64_t MERGE_LIST_IDX_THREE = 3;
|
||||
|
||||
__aicore__ inline int64_t Ceil(int64_t a, int64_t b)
|
||||
{
|
||||
if (b == 0) {
|
||||
return 0;
|
||||
}
|
||||
return (a + b - 1) / b;
|
||||
}
|
||||
|
||||
__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes)
|
||||
{
|
||||
if (bytes == 0) {
|
||||
return 0;
|
||||
}
|
||||
return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES / bytes;
|
||||
}
|
||||
|
||||
__aicore__ inline int64_t AlignBytes(int64_t elementNum, int64_t bytes)
|
||||
{
|
||||
return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline T Min(T a, T b)
|
||||
{
|
||||
return a > b ? b : a;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline T Max(T a, T b)
|
||||
{
|
||||
return a < b ? b : a;
|
||||
}
|
||||
|
||||
template <typename T1, typename T2>
|
||||
__aicore__ inline T1 CeilAlign(T1 a, T2 b) {
|
||||
if (b == 0) {
|
||||
return 0;
|
||||
}
|
||||
return (a + b - 1) / b * b;
|
||||
}
|
||||
|
||||
template <typename T1, typename T2>
|
||||
__aicore__ inline T1 CeilDiv(T1 x, T2 y)
|
||||
{
|
||||
if (y != 0 && x != 0) {
|
||||
const T1 quotient = x / y;
|
||||
return (x % y != 0 && ((x ^ y) >= 0)) ? (quotient + 1) : quotient;
|
||||
}
|
||||
|
||||
return x;
|
||||
}
|
||||
|
||||
template <HardEvent event>
|
||||
__aicore__ inline void SetWaitFlag(HardEvent evt)
|
||||
{
|
||||
event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(evt));
|
||||
SetFlag<event>(eventId);
|
||||
WaitFlag<event>(eventId);
|
||||
}
|
||||
|
||||
constexpr AscendC::MicroAPI::CastTrait castTraitB162B32 = {
|
||||
AscendC::MicroAPI::RegLayout::ZERO,
|
||||
AscendC::MicroAPI::SatMode::UNKNOWN,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::UNKNOWN,
|
||||
};
|
||||
|
||||
constexpr AscendC::MicroAPI::CastTrait castTraitB322B16 = {
|
||||
AscendC::MicroAPI::RegLayout::ZERO,
|
||||
AscendC::MicroAPI::SatMode::NO_SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_RINT,
|
||||
};
|
||||
|
||||
constexpr AscendC::MicroAPI::CastTrait castTraitB322Int32 = {
|
||||
AscendC::MicroAPI::RegLayout::UNKNOWN,
|
||||
AscendC::MicroAPI::SatMode::NO_SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_TRUNC,
|
||||
};
|
||||
|
||||
constexpr AscendC::MicroAPI::CastTrait castTraitB322Int16 = {
|
||||
AscendC::MicroAPI::RegLayout::ZERO,
|
||||
AscendC::MicroAPI::SatMode::NO_SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_TRUNC,
|
||||
};
|
||||
|
||||
constexpr AscendC::MicroAPI::CastTrait castTraitB162Int8 = {
|
||||
AscendC::MicroAPI::RegLayout::ZERO,
|
||||
AscendC::MicroAPI::SatMode::NO_SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_TRUNC,
|
||||
};
|
||||
|
||||
// load 对齐的 bfloat16,float16,bfloat32类型的 input(ub中)数据到 float32类型的dst(寄存器)中
|
||||
template <typename T>
|
||||
__aicore__ inline void LoadOneTensorForDtypeT(__local_mem__ T *input, MicroAPI::RegTensor<float> &dst,
|
||||
MicroAPI::MaskReg &preg, uint32_t offset)
|
||||
{
|
||||
if constexpr (IsSameType<T, half>::value) {
|
||||
MicroAPI::RegTensor<half> xFp16;
|
||||
DataCopy<half, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ half *)(input) + (offset)));
|
||||
Cast<float, half, castTraitB162B32>(dst, xFp16, preg);
|
||||
} else if constexpr (IsSameType<T, bfloat16_t>::value) {
|
||||
MicroAPI::RegTensor<bfloat16_t> xBf16;
|
||||
DataCopy<bfloat16_t, MicroAPI::LoadDist::DIST_UNPACK_B16>(xBf16,
|
||||
((__local_mem__ bfloat16_t *)(input) + (offset)));
|
||||
Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg);
|
||||
} else {
|
||||
DataCopy(dst, ((__local_mem__ float *)(input) + (offset)));
|
||||
}
|
||||
}
|
||||
|
||||
// load 2个对齐的Tensor 到寄存器中
|
||||
template <typename T>
|
||||
__aicore__ inline void LoadTwoTensorForDtypeT(__local_mem__ T *src1, __local_mem__ T *src2,
|
||||
MicroAPI::RegTensor<float> &dst1, MicroAPI::RegTensor<float> &dst2,
|
||||
MicroAPI::MaskReg &dst1Preg, MicroAPI::MaskReg &dst2Preg,
|
||||
uint32_t src1Offset, uint32_t src2Offset)
|
||||
{
|
||||
if constexpr (IsSameType<T, half>::value) {
|
||||
MicroAPI::RegTensor<half> xFp16Q;
|
||||
MicroAPI::RegTensor<half> xFp16R;
|
||||
DataCopy<half, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16Q, ((__local_mem__ half *)(src1) + (src1Offset)));
|
||||
DataCopy<half, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16R, ((__local_mem__ half *)(src2) + (src2Offset)));
|
||||
Cast<float, half, castTraitB162B32>(dst1, xFp16Q, dst1Preg);
|
||||
Cast<float, half, castTraitB162B32>(dst2, xFp16R, dst2Preg);
|
||||
} else if constexpr (IsSameType<T, bfloat16_t>::value) {
|
||||
MicroAPI::RegTensor<bfloat16_t> xFp16Q;
|
||||
MicroAPI::RegTensor<bfloat16_t> xFp16R;
|
||||
DataCopy<bfloat16_t, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16Q, ((__local_mem__ bfloat16_t *)(src1) + (src1Offset)));
|
||||
DataCopy<bfloat16_t, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16R, ((__local_mem__ bfloat16_t *)(src2) + (src2Offset)));
|
||||
Cast<float, bfloat16_t, castTraitB162B32>(dst1, xFp16Q, dst1Preg);
|
||||
Cast<float, bfloat16_t, castTraitB162B32>(dst2, xFp16R, dst2Preg);
|
||||
} else {
|
||||
DataCopy(dst1, ((__local_mem__ float *)(src1) + (src1Offset)));
|
||||
DataCopy(dst2, ((__local_mem__ float *)(src2) + (src2Offset)));
|
||||
}
|
||||
}
|
||||
|
||||
// store 对齐的float32类型的src(寄存器)数据到output(ub)中，output数据类型支持bfloat16,float16,bfloat32,int32_t,int16_t,int8_t,uint8_t
|
||||
template <typename T>
|
||||
__aicore__ inline void StoreOneTensorForDtypeT(__local_mem__ T *output, MicroAPI::RegTensor<float> &src,
|
||||
MicroAPI::MaskReg &preg, uint32_t offset)
|
||||
{
|
||||
if constexpr (IsSameType<T, half>::value) {
|
||||
MicroAPI::RegTensor<half> yFp16;
|
||||
Cast<half, float, castTraitB322B16>(yFp16, src, preg);
|
||||
DataCopy<half, MicroAPI::StoreDist::DIST_PACK_B32>(((__local_mem__ half *)output + offset), yFp16, preg);
|
||||
} else if constexpr (IsSameType<T, bfloat16_t>::value) {
|
||||
MicroAPI::RegTensor<bfloat16_t> xBf16;
|
||||
Cast<bfloat16_t, float, castTraitB322B16>(xBf16, src, preg);
|
||||
DataCopy<bfloat16_t, MicroAPI::StoreDist::DIST_PACK_B32>(((__local_mem__ bfloat16_t *)output + offset),
|
||||
xBf16, preg);
|
||||
} else if constexpr (IsSameType<T, int32_t>::value) {
|
||||
MicroAPI::RegTensor<int32_t> zInt32;
|
||||
MicroAPI::Cast<int32_t, float, castTraitB322Int32>(zInt32, src, preg);
|
||||
DataCopy<int32_t, AscendC::MicroAPI::StoreDist::DIST_NORM>(((__local_mem__ int32_t *)output + offset), zInt32, preg);
|
||||
} else if constexpr (IsSameType<T, int16_t>::value) {
|
||||
MicroAPI::RegTensor<int16_t> zInt16;
|
||||
MicroAPI::Cast<int16_t, float, castTraitB322Int16>(zInt16, src, preg);
|
||||
DataCopy<int16_t, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(((__local_mem__ int16_t *)output + offset), zInt16, preg);
|
||||
} else if constexpr (IsSameType<T, int8_t>::value) {
|
||||
MicroAPI::RegTensor<half> yFp16;
|
||||
MicroAPI::RegTensor<int8_t> zInt8;
|
||||
Cast<half, float, castTraitB322Int16>(yFp16, src, preg);
|
||||
Cast<int8_t, half, castTraitB162Int8>(zInt8, yFp16, preg);
|
||||
DataCopy<int8_t, MicroAPI::StoreDist::DIST_PACK4_B32>(((__local_mem__ int8_t *)output + offset), zInt8, preg);
|
||||
} else if constexpr (IsSameType<T, uint8_t>::value) {
|
||||
MicroAPI::RegTensor<half> yFp16;
|
||||
MicroAPI::RegTensor<uint8_t> zUint8;
|
||||
Cast<half, float, castTraitB322Int16>(yFp16, src, preg);
|
||||
Cast<uint8_t, half, castTraitB162Int8>(zUint8, yFp16, preg);
|
||||
DataCopy<uint8_t, MicroAPI::StoreDist::DIST_PACK4_B32>(((__local_mem__ uint8_t *)output + offset), zUint8, preg);
|
||||
} else {
|
||||
DataCopy(((__local_mem__ float *)output + offset), src, preg);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace MoeGatingTopKHash
|
||||
#endif // MOE_GATING_TOP_K_COMMON_H_V35
|
||||
File diff suppressed because it is too large
Load Diff
90
csrc/moe/moe_gating_top_k_hash/op_kernel/common.h
Normal file
90
csrc/moe/moe_gating_top_k_hash/op_kernel/common.h
Normal file
@@ -0,0 +1,90 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file common.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef MOE_GATING_TOP_K_COMMON_H
|
||||
#define MOE_GATING_TOP_K_COMMON_H
|
||||
|
||||
#include "kernel_operator.h"
|
||||
|
||||
namespace MoeGatingTopKHash {
|
||||
using namespace AscendC;
|
||||
const float MIN_FP32 = *(float *)(&F32_NEG_INF);
|
||||
constexpr int32_t FLOAT32_NEG_INF = 0xFF800000; // -inf -2139095040
|
||||
constexpr int64_t ONE_REPEAT_SORT_NUM = 32;
|
||||
constexpr int64_t BLOCK_BYTES = 32;
|
||||
constexpr int64_t REPEAT_BYTES = 256;
|
||||
constexpr int64_t REPEAT_BLOCKS = 8;
|
||||
|
||||
constexpr int32_t CONSTANT_TWO = 2;
|
||||
constexpr int32_t CONSTANT_THREE = 3;
|
||||
constexpr int32_t CONSTANT_FOUR = 4;
|
||||
constexpr int32_t CONSTANT_EIGHT = 8;
|
||||
|
||||
constexpr int64_t MERGE_LIST_TWO = 2;
|
||||
constexpr int64_t MERGE_LIST_THREE = 3;
|
||||
constexpr int64_t MERGE_LIST_FOUR = 4;
|
||||
|
||||
constexpr int64_t MERGE_LIST_IDX_TWO = 2;
|
||||
constexpr int64_t MERGE_LIST_IDX_THREE = 3;
|
||||
|
||||
constexpr int64_t NORM_TYPE_SOFTMAX = 0;
|
||||
constexpr int64_t NORM_TYPE_SIGMOID = 1;
|
||||
constexpr int64_t NORM_TYPE_SOFTPLUS = 2;
|
||||
|
||||
__aicore__ inline int64_t Ceil(int64_t a, int64_t b)
|
||||
{
|
||||
if (b == 0) {
|
||||
return 0;
|
||||
}
|
||||
return (a + b - 1) / b;
|
||||
}
|
||||
|
||||
__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes)
|
||||
{
|
||||
if (bytes == 0) {
|
||||
return 0;
|
||||
}
|
||||
return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES / bytes;
|
||||
}
|
||||
|
||||
__aicore__ inline int64_t AlignBytes(int64_t elementNum, int64_t bytes)
|
||||
{
|
||||
return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline T Min(T a, T b)
|
||||
{
|
||||
return a > b ? b : a;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline T Max(T a, T b)
|
||||
{
|
||||
return a < b ? b : a;
|
||||
}
|
||||
|
||||
template <typename T1, typename T2>
|
||||
__aicore__ inline T1 CeilDiv(T1 x, T2 y)
|
||||
{
|
||||
if (y != 0 && x != 0) {
|
||||
const T1 quotient = x / y;
|
||||
return (x % y != 0 && ((x ^ y) >= 0)) ? (quotient + 1) : quotient;
|
||||
}
|
||||
|
||||
return x;
|
||||
}
|
||||
|
||||
} // namespace MoeGatingTopKHash
|
||||
#endif // MOE_GATING_TOP_K_COMMON_H
|
||||
214
csrc/moe/moe_gating_top_k_hash/op_kernel/common_regbase.h
Normal file
214
csrc/moe/moe_gating_top_k_hash/op_kernel/common_regbase.h
Normal file
@@ -0,0 +1,214 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file common.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef MOE_GATING_TOP_K_COMMON_H_V35
|
||||
#define MOE_GATING_TOP_K_COMMON_H_V35
|
||||
|
||||
#include "kernel_operator.h"
|
||||
|
||||
namespace MoeGatingTopKHashRegbaseNS {
|
||||
using namespace AscendC;
|
||||
constexpr int32_t MIN_FP32 = 0xFF800000;
|
||||
constexpr int64_t ONE_REPEAT_SORT_NUM = 32;
|
||||
constexpr int64_t BLOCK_BYTES = 32;
|
||||
|
||||
constexpr int64_t MERGE_LIST_TWO = 2;
|
||||
constexpr int64_t MERGE_LIST_THREE = 3;
|
||||
constexpr int64_t MERGE_LIST_FOUR = 4;
|
||||
|
||||
constexpr int64_t MERGE_LIST_IDX_TWO = 2;
|
||||
constexpr int64_t MERGE_LIST_IDX_THREE = 3;
|
||||
|
||||
__aicore__ inline int64_t Ceil(int64_t a, int64_t b)
|
||||
{
|
||||
if (b == 0) {
|
||||
return 0;
|
||||
}
|
||||
return (a + b - 1) / b;
|
||||
}
|
||||
|
||||
__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes)
|
||||
{
|
||||
if (bytes == 0) {
|
||||
return 0;
|
||||
}
|
||||
return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES / bytes;
|
||||
}
|
||||
|
||||
__aicore__ inline int64_t AlignBytes(int64_t elementNum, int64_t bytes)
|
||||
{
|
||||
return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline T Min(T a, T b)
|
||||
{
|
||||
return a > b ? b : a;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline T Max(T a, T b)
|
||||
{
|
||||
return a < b ? b : a;
|
||||
}
|
||||
|
||||
template <typename T1, typename T2>
|
||||
__aicore__ inline T1 CeilAlign(T1 a, T2 b) {
|
||||
if (b == 0) {
|
||||
return 0;
|
||||
}
|
||||
return (a + b - 1) / b * b;
|
||||
}
|
||||
|
||||
template <typename T1, typename T2>
|
||||
__aicore__ inline T1 CeilDiv(T1 x, T2 y)
|
||||
{
|
||||
if (y != 0 && x != 0) {
|
||||
const T1 quotient = x / y;
|
||||
return (x % y != 0 && ((x ^ y) >= 0)) ? (quotient + 1) : quotient;
|
||||
}
|
||||
|
||||
return x;
|
||||
}
|
||||
|
||||
template <HardEvent event>
|
||||
__aicore__ inline void SetWaitFlag(HardEvent evt)
|
||||
{
|
||||
event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(evt));
|
||||
SetFlag<event>(eventId);
|
||||
WaitFlag<event>(eventId);
|
||||
}
|
||||
|
||||
constexpr AscendC::MicroAPI::CastTrait castTraitB162B32 = {
|
||||
AscendC::MicroAPI::RegLayout::ZERO,
|
||||
AscendC::MicroAPI::SatMode::UNKNOWN,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::UNKNOWN,
|
||||
};
|
||||
|
||||
constexpr AscendC::MicroAPI::CastTrait castTraitB322B16 = {
|
||||
AscendC::MicroAPI::RegLayout::ZERO,
|
||||
AscendC::MicroAPI::SatMode::NO_SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_RINT,
|
||||
};
|
||||
|
||||
constexpr AscendC::MicroAPI::CastTrait castTraitB322Int32 = {
|
||||
AscendC::MicroAPI::RegLayout::UNKNOWN,
|
||||
AscendC::MicroAPI::SatMode::NO_SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_TRUNC,
|
||||
};
|
||||
|
||||
constexpr AscendC::MicroAPI::CastTrait castTraitB322Int16 = {
|
||||
AscendC::MicroAPI::RegLayout::ZERO,
|
||||
AscendC::MicroAPI::SatMode::NO_SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_TRUNC,
|
||||
};
|
||||
|
||||
constexpr AscendC::MicroAPI::CastTrait castTraitB162Int8 = {
|
||||
AscendC::MicroAPI::RegLayout::ZERO,
|
||||
AscendC::MicroAPI::SatMode::NO_SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_TRUNC,
|
||||
};
|
||||
|
||||
// load 对齐的 bfloat16,float16,bfloat32类型的 input(ub中)数据到 float32类型的dst(寄存器)中
|
||||
template <typename T>
|
||||
__aicore__ inline void LoadOneTensorForDtypeT(__local_mem__ T *input, MicroAPI::RegTensor<float> &dst,
|
||||
MicroAPI::MaskReg &preg, uint32_t offset)
|
||||
{
|
||||
if constexpr (IsSameType<T, half>::value) {
|
||||
MicroAPI::RegTensor<half> xFp16;
|
||||
DataCopy<half, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ half *)(input) + (offset)));
|
||||
Cast<float, half, castTraitB162B32>(dst, xFp16, preg);
|
||||
} else if constexpr (IsSameType<T, bfloat16_t>::value) {
|
||||
MicroAPI::RegTensor<bfloat16_t> xBf16;
|
||||
DataCopy<bfloat16_t, MicroAPI::LoadDist::DIST_UNPACK_B16>(xBf16,
|
||||
((__local_mem__ bfloat16_t *)(input) + (offset)));
|
||||
Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg);
|
||||
} else {
|
||||
DataCopy(dst, ((__local_mem__ float *)(input) + (offset)));
|
||||
}
|
||||
}
|
||||
|
||||
// load 2个对齐的Tensor 到寄存器中
|
||||
template <typename T>
|
||||
__aicore__ inline void LoadTwoTensorForDtypeT(__local_mem__ T *src1, __local_mem__ T *src2,
|
||||
MicroAPI::RegTensor<float> &dst1, MicroAPI::RegTensor<float> &dst2,
|
||||
MicroAPI::MaskReg &dst1Preg, MicroAPI::MaskReg &dst2Preg,
|
||||
uint32_t src1Offset, uint32_t src2Offset)
|
||||
{
|
||||
if constexpr (IsSameType<T, half>::value) {
|
||||
MicroAPI::RegTensor<half> xFp16Q;
|
||||
MicroAPI::RegTensor<half> xFp16R;
|
||||
DataCopy<half, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16Q, ((__local_mem__ half *)(src1) + (src1Offset)));
|
||||
DataCopy<half, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16R, ((__local_mem__ half *)(src2) + (src2Offset)));
|
||||
Cast<float, half, castTraitB162B32>(dst1, xFp16Q, dst1Preg);
|
||||
Cast<float, half, castTraitB162B32>(dst2, xFp16R, dst2Preg);
|
||||
} else if constexpr (IsSameType<T, bfloat16_t>::value) {
|
||||
MicroAPI::RegTensor<bfloat16_t> xFp16Q;
|
||||
MicroAPI::RegTensor<bfloat16_t> xFp16R;
|
||||
DataCopy<bfloat16_t, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16Q, ((__local_mem__ bfloat16_t *)(src1) + (src1Offset)));
|
||||
DataCopy<bfloat16_t, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16R, ((__local_mem__ bfloat16_t *)(src2) + (src2Offset)));
|
||||
Cast<float, bfloat16_t, castTraitB162B32>(dst1, xFp16Q, dst1Preg);
|
||||
Cast<float, bfloat16_t, castTraitB162B32>(dst2, xFp16R, dst2Preg);
|
||||
} else {
|
||||
DataCopy(dst1, ((__local_mem__ float *)(src1) + (src1Offset)));
|
||||
DataCopy(dst2, ((__local_mem__ float *)(src2) + (src2Offset)));
|
||||
}
|
||||
}
|
||||
|
||||
// store 对齐的float32类型的src(寄存器)数据到output(ub)中，output数据类型支持bfloat16,float16,bfloat32,int32_t,int16_t,int8_t,uint8_t
|
||||
template <typename T>
|
||||
__aicore__ inline void StoreOneTensorForDtypeT(__local_mem__ T *output, MicroAPI::RegTensor<float> &src,
|
||||
MicroAPI::MaskReg &preg, uint32_t offset)
|
||||
{
|
||||
if constexpr (IsSameType<T, half>::value) {
|
||||
MicroAPI::RegTensor<half> yFp16;
|
||||
Cast<half, float, castTraitB322B16>(yFp16, src, preg);
|
||||
DataCopy<half, MicroAPI::StoreDist::DIST_PACK_B32>(((__local_mem__ half *)output + offset), yFp16, preg);
|
||||
} else if constexpr (IsSameType<T, bfloat16_t>::value) {
|
||||
MicroAPI::RegTensor<bfloat16_t> xBf16;
|
||||
Cast<bfloat16_t, float, castTraitB322B16>(xBf16, src, preg);
|
||||
DataCopy<bfloat16_t, MicroAPI::StoreDist::DIST_PACK_B32>(((__local_mem__ bfloat16_t *)output + offset),
|
||||
xBf16, preg);
|
||||
} else if constexpr (IsSameType<T, int32_t>::value) {
|
||||
MicroAPI::RegTensor<int32_t> zInt32;
|
||||
MicroAPI::Cast<int32_t, float, castTraitB322Int32>(zInt32, src, preg);
|
||||
DataCopy<int32_t, AscendC::MicroAPI::StoreDist::DIST_NORM>(((__local_mem__ int32_t *)output + offset), zInt32, preg);
|
||||
} else if constexpr (IsSameType<T, int16_t>::value) {
|
||||
MicroAPI::RegTensor<int16_t> zInt16;
|
||||
MicroAPI::Cast<int16_t, float, castTraitB322Int16>(zInt16, src, preg);
|
||||
DataCopy<int16_t, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(((__local_mem__ int16_t *)output + offset), zInt16, preg);
|
||||
} else if constexpr (IsSameType<T, int8_t>::value) {
|
||||
MicroAPI::RegTensor<half> yFp16;
|
||||
MicroAPI::RegTensor<int8_t> zInt8;
|
||||
Cast<half, float, castTraitB322Int16>(yFp16, src, preg);
|
||||
Cast<int8_t, half, castTraitB162Int8>(zInt8, yFp16, preg);
|
||||
DataCopy<int8_t, MicroAPI::StoreDist::DIST_PACK4_B32>(((__local_mem__ int8_t *)output + offset), zInt8, preg);
|
||||
} else if constexpr (IsSameType<T, uint8_t>::value) {
|
||||
MicroAPI::RegTensor<half> yFp16;
|
||||
MicroAPI::RegTensor<uint8_t> zUint8;
|
||||
Cast<half, float, castTraitB322Int16>(yFp16, src, preg);
|
||||
Cast<uint8_t, half, castTraitB162Int8>(zUint8, yFp16, preg);
|
||||
DataCopy<uint8_t, MicroAPI::StoreDist::DIST_PACK4_B32>(((__local_mem__ uint8_t *)output + offset), zUint8, preg);
|
||||
} else {
|
||||
DataCopy(((__local_mem__ float *)output + offset), src, preg);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace MoeGatingTopKHash
|
||||
#endif // MOE_GATING_TOP_K_COMMON_H_V35
|
||||
@@ -0,0 +1,132 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file moe_gating_top_k_hash.cpp
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#include "moe_gating_top_k_hash_e_k_fullload.h"
|
||||
#include "moe_gating_top_k_hash_without_group.h"
|
||||
#include "moe_gating_top_k_hash_generalized.h"
|
||||
#if defined(__DAV_C310__)
|
||||
#include "moe_gating_top_k_hash_regbase.h"
|
||||
using namespace MoeGatingTopKHashRegbaseNS;
|
||||
#endif
|
||||
#define TILING_KEY_PER_GROUP_COUNT_32 0
|
||||
#define TILING_KEY_WITHOUT_GROUP 1
|
||||
#define TILING_KEY_GENERALIZED 2
|
||||
#define TILING_KEY_WITHOUT_GROUP_1 3
|
||||
#define TILING_KEY_WITHOUT_GROUP_2 4
|
||||
#define TILING_KEY_WITHOUT_GROUP_3 5
|
||||
#define TILING_KEY_WITHOUT_GROUP_4 6
|
||||
#define TILING_KEY_REGBASE 10000
|
||||
#define TILING_KEY_REGBASE_1 10001
|
||||
#define TILING_KEY_REGBASE_2 10002
|
||||
#define TILING_KEY_REGBASE_3 10003
|
||||
#define TILING_KEY_REGBASE_4 10004
|
||||
|
||||
using namespace AscendC;
|
||||
using namespace MoeGatingTopKHash;
|
||||
extern "C" __global__ __aicore__ void moe_gating_top_k_hash(GM_ADDR x, GM_ADDR bias, GM_ADDR inputIds, GM_ADDR tid2eid, GM_ADDR y, GM_ADDR expertIdx,
|
||||
GM_ADDR out, GM_ADDR workspace, GM_ADDR tiling)
|
||||
{
|
||||
KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
|
||||
if (g_coreType == AIC) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (workspace == nullptr) {
|
||||
return;
|
||||
}
|
||||
|
||||
GM_ADDR userWS = GetUserWorkspace(workspace);
|
||||
if (userWS == nullptr) {
|
||||
return;
|
||||
}
|
||||
|
||||
TPipe tPipe;
|
||||
if (TILING_KEY_IS(TILING_KEY_PER_GROUP_COUNT_32)) {
|
||||
GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashTilingData, tilingData, tiling);
|
||||
const MoeGatingTopKHashTilingData *__restrict t = &tilingData;
|
||||
MoeGatingTopKHashEKFullload<DTYPE_X> op;
|
||||
op.Init(x, bias, y, expertIdx, out, userWS, t, &tPipe);
|
||||
op.Process();
|
||||
} else if (TILING_KEY_IS(TILING_KEY_WITHOUT_GROUP)) {
|
||||
GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashTilingData, tilingData, tiling);
|
||||
const MoeGatingTopKHashTilingData *__restrict t = &tilingData;
|
||||
MoeGatingTopKHashWithoutGroup<DTYPE_X, int32_t, int32_t> op;
|
||||
op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, t, &tPipe);
|
||||
op.Process();
|
||||
} else if (TILING_KEY_IS(TILING_KEY_WITHOUT_GROUP_1)) {
|
||||
GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashTilingData, tilingData, tiling);
|
||||
const MoeGatingTopKHashTilingData *__restrict t = &tilingData;
|
||||
MoeGatingTopKHashWithoutGroup<DTYPE_X, int32_t, int64_t> op;
|
||||
op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, t, &tPipe);
|
||||
op.Process();
|
||||
} else if (TILING_KEY_IS(TILING_KEY_WITHOUT_GROUP_2)) {
|
||||
GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashTilingData, tilingData, tiling);
|
||||
const MoeGatingTopKHashTilingData *__restrict t = &tilingData;
|
||||
MoeGatingTopKHashWithoutGroup<DTYPE_X, int32_t, int32_t> op;
|
||||
op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, t, &tPipe);
|
||||
op.Process();
|
||||
} else if (TILING_KEY_IS(TILING_KEY_WITHOUT_GROUP_3)) {
|
||||
GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashTilingData, tilingData, tiling);
|
||||
const MoeGatingTopKHashTilingData *__restrict t = &tilingData;
|
||||
MoeGatingTopKHashWithoutGroup<DTYPE_X, int64_t, int64_t> op;
|
||||
op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, t, &tPipe);
|
||||
op.Process();
|
||||
} else if (TILING_KEY_IS(TILING_KEY_WITHOUT_GROUP_4)) {
|
||||
GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashTilingData, tilingData, tiling);
|
||||
const MoeGatingTopKHashTilingData *__restrict t = &tilingData;
|
||||
MoeGatingTopKHashWithoutGroup<DTYPE_X, int64_t, int32_t> op;
|
||||
op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, t, &tPipe);
|
||||
op.Process();
|
||||
} else if (TILING_KEY_IS(TILING_KEY_GENERALIZED)) {
|
||||
GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashTilingData, tilingData, tiling);
|
||||
const MoeGatingTopKHashTilingData *__restrict t = &tilingData;
|
||||
MoeGatingTopKHashGenerlized<DTYPE_X> op;
|
||||
op.Init(x, bias, y, expertIdx, out, userWS, t, &tPipe);
|
||||
op.Process();
|
||||
}
|
||||
#if defined(__DAV_C310__)
|
||||
else if (TILING_KEY_IS(TILING_KEY_REGBASE)) {
|
||||
GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashRegbaseTilingData, tiling_data_in, tiling);
|
||||
const MoeGatingTopKHashRegbaseTilingData *__restrict tilingData = &tiling_data_in;
|
||||
MoeGatingTopKHashRegbase<DTYPE_X, int32_t, int32_t> op;
|
||||
op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe);
|
||||
op.Process();
|
||||
} else if (TILING_KEY_IS(TILING_KEY_REGBASE_1)) {
|
||||
GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashRegbaseTilingData, tiling_data_in, tiling);
|
||||
const MoeGatingTopKHashRegbaseTilingData *__restrict tilingData = &tiling_data_in;
|
||||
MoeGatingTopKHashRegbase<DTYPE_X, int32_t, int64_t> op;
|
||||
op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe);
|
||||
op.Process();
|
||||
} else if (TILING_KEY_IS(TILING_KEY_REGBASE_2)) {
|
||||
GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashRegbaseTilingData, tiling_data_in, tiling);
|
||||
const MoeGatingTopKHashRegbaseTilingData *__restrict tilingData = &tiling_data_in;
|
||||
MoeGatingTopKHashRegbase<DTYPE_X, int32_t, int32_t> op;
|
||||
op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe);
|
||||
op.Process();
|
||||
} else if (TILING_KEY_IS(TILING_KEY_REGBASE_3)) {
|
||||
GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashRegbaseTilingData, tiling_data_in, tiling);
|
||||
const MoeGatingTopKHashRegbaseTilingData *__restrict tilingData = &tiling_data_in;
|
||||
MoeGatingTopKHashRegbase<DTYPE_X, int64_t, int64_t> op;
|
||||
op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe);
|
||||
op.Process();
|
||||
} else if (TILING_KEY_IS(TILING_KEY_REGBASE_4)) {
|
||||
GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashRegbaseTilingData, tiling_data_in, tiling);
|
||||
const MoeGatingTopKHashRegbaseTilingData *__restrict tilingData = &tiling_data_in;
|
||||
MoeGatingTopKHashRegbase<DTYPE_X, int64_t, int32_t> op;
|
||||
op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe);
|
||||
op.Process();
|
||||
}
|
||||
#endif
|
||||
}
|
||||
@@ -0,0 +1,404 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file moe_gating_top_k_hash_e_k_fullload.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef MOE_GATING_TOP_K_E_K_FULLLOAD_H
|
||||
#define MOE_GATING_TOP_K_E_K_FULLLOAD_H
|
||||
#include "kernel_operator.h"
|
||||
#include "common.h"
|
||||
namespace MoeGatingTopKHash {
|
||||
using namespace AscendC;
|
||||
|
||||
template <typename T>
|
||||
class MoeGatingTopKHashEKFullload {
|
||||
public:
|
||||
__aicore__ inline MoeGatingTopKHashEKFullload(){};
|
||||
__aicore__ inline void Init(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, GM_ADDR workspace,
|
||||
const MoeGatingTopKHashTilingData *tilingData, TPipe *tPipe);
|
||||
__aicore__ inline void Process();
|
||||
|
||||
private:
|
||||
__aicore__ inline void CopyInBias();
|
||||
__aicore__ inline void CopyInX(int64_t progress);
|
||||
__aicore__ inline void ComputeX();
|
||||
__aicore__ inline void SortInGroup();
|
||||
__aicore__ inline void SelectTopKGroupIndex();
|
||||
__aicore__ inline void SelectTopKExpertIdx();
|
||||
__aicore__ inline void SelectTopKExpertScore();
|
||||
__aicore__ inline void CopyOut(int64_t progress);
|
||||
|
||||
private:
|
||||
TPipe *pipe_;
|
||||
TQue<QuePosition::VECIN, 1> xInQueue_;
|
||||
TBuf<TPosition::VECCALC> biasInQueue_;
|
||||
TQue<QuePosition::VECOUT, 1> yOutQueue_;
|
||||
TQue<QuePosition::VECOUT, 1> expertIdxOutQueue_;
|
||||
TQue<QuePosition::VECOUT, 1> outOutQueue_;
|
||||
|
||||
TQue<QuePosition::VECOUT, 1> xBiasQueue_;
|
||||
TQue<QuePosition::VECOUT, 1> xSigmoidQueue_;
|
||||
TQue<QuePosition::VECIN, 1> sigmoidTmpQueue_;
|
||||
TQue<QuePosition::VECIN, 1> sortedInGroupQueue_;
|
||||
TQue<QuePosition::VECIN, 1> sortedGroupQueue_;
|
||||
TBuf<TPosition::VECCALC> calcTmpBuffer_;
|
||||
|
||||
GlobalTensor<T> xGm_;
|
||||
GlobalTensor<T> biasGm_;
|
||||
GlobalTensor<T> yGm_;
|
||||
GlobalTensor<int32_t> expertIdxGm_;
|
||||
GlobalTensor<T> outGm_;
|
||||
|
||||
int64_t blockIdx_;
|
||||
int64_t perCoreRowCount_;
|
||||
int64_t curCoreRowCount_;
|
||||
int64_t expertCount_;
|
||||
bool addBias_;
|
||||
int64_t k_;
|
||||
int64_t kGroup_;
|
||||
int64_t groupCount_;
|
||||
int64_t groupSelectMode_;
|
||||
int64_t renorm_;
|
||||
int64_t normType_;
|
||||
int64_t outFlag_;
|
||||
float routedScalingFactor_;
|
||||
float eps_;
|
||||
|
||||
int64_t expertCountAlign_;
|
||||
int64_t kAlign_;
|
||||
int64_t perGroupExpertCount_;
|
||||
|
||||
const MoeGatingTopKHashTilingData *tilingData_;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashEKFullload<T>::CopyInBias()
|
||||
{
|
||||
LocalTensor<float> biasTensor = biasInQueue_.Get<float>();
|
||||
DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(expertCount_ * sizeof(T)), 0, 0, 0};
|
||||
DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast<T>(0)};
|
||||
if constexpr (IsSameType<T, float>::value) {
|
||||
DataCopyPad(biasTensor, biasGm_, dataCopyParams, dataCopyPadParams);
|
||||
event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
|
||||
SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
|
||||
WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
|
||||
} else {
|
||||
DataCopyPad(biasTensor[expertCountAlign_].ReinterpretCast<T>(), biasGm_, dataCopyParams, dataCopyPadParams);
|
||||
event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
|
||||
SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
|
||||
WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
|
||||
Cast(biasTensor, biasTensor[expertCountAlign_].ReinterpretCast<T>(), RoundMode::CAST_NONE, expertCount_);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashEKFullload<T>::CopyInX(int64_t row)
|
||||
{
|
||||
LocalTensor<float> xInLocalTensor = xInQueue_.AllocTensor<float>();
|
||||
DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(expertCount_ * sizeof(T)), 0, 0, 0};
|
||||
DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast<T>(0)};
|
||||
if constexpr (IsSameType<T, float>::value) {
|
||||
DataCopyPad(xInLocalTensor, xGm_[row * expertCount_], dataCopyParams, dataCopyPadParams);
|
||||
} else {
|
||||
DataCopyPad(xInLocalTensor[expertCountAlign_].ReinterpretCast<T>(), xGm_[row * expertCount_], dataCopyParams,
|
||||
dataCopyPadParams);
|
||||
event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
|
||||
SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
|
||||
WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
|
||||
Cast(xInLocalTensor, xInLocalTensor[expertCountAlign_].ReinterpretCast<T>(), RoundMode::CAST_NONE,
|
||||
expertCount_);
|
||||
}
|
||||
|
||||
xInQueue_.EnQue(xInLocalTensor);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashEKFullload<T>::ComputeX()
|
||||
{
|
||||
LocalTensor<float> xSigmoidTensor = xSigmoidQueue_.AllocTensor<float>();
|
||||
LocalTensor<float> xInLocalTensor = xInQueue_.DeQue<float>();
|
||||
LocalTensor<float> xBiasTensor = xBiasQueue_.AllocTensor<float>();
|
||||
LocalTensor<float> biasTensor = biasInQueue_.Get<float>();
|
||||
LocalTensor<uint8_t> sharedTmpBuffer = sigmoidTmpQueue_.AllocTensor<uint8_t>(); // 临时空间可以复用
|
||||
Sigmoid(xSigmoidTensor, xInLocalTensor, sharedTmpBuffer, expertCount_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
if (addBias_) {
|
||||
Add(xBiasTensor, xSigmoidTensor, biasTensor, expertCount_);
|
||||
} else {
|
||||
Adds(xBiasTensor, xSigmoidTensor, static_cast<float>(0), expertCount_);
|
||||
}
|
||||
|
||||
xSigmoidQueue_.EnQue<float>(xSigmoidTensor);
|
||||
xBiasQueue_.EnQue<float>(xBiasTensor);
|
||||
xInQueue_.FreeTensor(xInLocalTensor);
|
||||
sigmoidTmpQueue_.FreeTensor(sharedTmpBuffer);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashEKFullload<T>::SortInGroup()
|
||||
{
|
||||
LocalTensor<float> xBiasTensor = xBiasQueue_.DeQue<float>();
|
||||
LocalTensor<float> sortedInGroupTensor = sortedInGroupQueue_.AllocTensor<float>(); // 组内排序的结果, 后续归并需要
|
||||
LocalTensor<uint32_t> indexTensor = calcTmpBuffer_.Get<uint32_t>(); // 用于存储排序时的索引
|
||||
ArithProgression(indexTensor.ReinterpretCast<int32_t>(), 0, 1, expertCount_); // 生成组索引0 1 2 ......
|
||||
PipeBarrier<PIPE_V>();
|
||||
Sort32(sortedInGroupTensor, xBiasTensor, indexTensor, expertCount_ / ONE_REPEAT_SORT_NUM); // 组内排序
|
||||
sortedInGroupQueue_.EnQue<float>(sortedInGroupTensor);
|
||||
xBiasQueue_.FreeTensor(xBiasTensor);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashEKFullload<T>::SelectTopKGroupIndex()
|
||||
{
|
||||
LocalTensor<float> sortedInGroupTensor = sortedInGroupQueue_.DeQue<float>();
|
||||
LocalTensor<uint32_t> indexTensor = calcTmpBuffer_.Get<uint32_t>();
|
||||
LocalTensor<float> top2ValueInGroupTensor = sigmoidTmpQueue_.AllocTensor<float>(); // 这个临时空间可以复用
|
||||
event_t eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
|
||||
indexTensor.SetValue(0, static_cast<uint32_t>(5)); // b0101
|
||||
indexTensor.SetValue(1, static_cast<uint32_t>(0));
|
||||
event_t eventIdSToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
|
||||
uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数
|
||||
GatherMaskParams gatherMaskParams;
|
||||
gatherMaskParams.repeatTimes = 8;
|
||||
gatherMaskParams.src0BlockStride = 1;
|
||||
gatherMaskParams.src0RepeatStride = 8;
|
||||
gatherMaskParams.src1RepeatStride = 0;
|
||||
GatherMask(top2ValueInGroupTensor, sortedInGroupTensor, indexTensor, true, static_cast<uint32_t>(64),
|
||||
gatherMaskParams, rsvdCnt);
|
||||
PipeBarrier<PIPE_V>();
|
||||
LocalTensor<float> groupTop2SumTensor = top2ValueInGroupTensor;
|
||||
PairReduceSum(groupTop2SumTensor, top2ValueInGroupTensor, 1, groupCount_ * 2, 1, 1,
|
||||
1); // 计算每个组内最大的两个数之和
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
LocalTensor<uint32_t> groupIndexTensor = indexTensor;
|
||||
ArithProgression(groupIndexTensor.ReinterpretCast<int32_t>(), 0, 1, groupCount_); // 生成组索引
|
||||
PipeBarrier<PIPE_V>();
|
||||
// 用最小值补到32个数
|
||||
int64_t duplicateNum = ONE_REPEAT_SORT_NUM - groupCount_;
|
||||
if (duplicateNum > 0) {
|
||||
uint64_t mask0 = UINT64_MAX << groupCount_;
|
||||
uint64_t mask[2] = {mask0, 0};
|
||||
Duplicate(groupTop2SumTensor, MIN_FP32, mask, 1, 1, 8);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
// 排序，将kgroup选出来
|
||||
LocalTensor<float> sortedGroupTensor = sortedGroupQueue_.AllocTensor<float>();
|
||||
Sort32(sortedGroupTensor, groupTop2SumTensor, groupIndexTensor, 1);
|
||||
|
||||
PipeBarrier<PIPE_V>();
|
||||
LocalTensor<int32_t> sortedGroupIndexTensor = indexTensor.ReinterpretCast<int32_t>();
|
||||
// 提取组序号
|
||||
uint8_t src1Pattern = 2; // 内置固定模式
|
||||
GatherMask(sortedGroupIndexTensor, sortedGroupTensor.template ReinterpretCast<int32_t>(), src1Pattern, false,
|
||||
static_cast<uint32_t>(0), {1, 1, 0, 0}, rsvdCnt);
|
||||
|
||||
// 需要将组排序(这里是降序，所以下mrgsor的时候反着取，3、2、1、0)
|
||||
Cast(sortedGroupTensor, sortedGroupIndexTensor, RoundMode::CAST_ROUND, kGroup_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
duplicateNum = ONE_REPEAT_SORT_NUM - kGroup_;
|
||||
if (duplicateNum > 0) {
|
||||
uint64_t mask0 = UINT64_MAX << kGroup_;
|
||||
uint64_t mask[2] = {mask0, 0};
|
||||
Duplicate(sortedGroupTensor, MIN_FP32, mask, 1, 1, 8);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
Sort32(top2ValueInGroupTensor, sortedGroupTensor, sortedGroupIndexTensor.template ReinterpretCast<uint32_t>(), 1);
|
||||
PipeBarrier<PIPE_V>();
|
||||
src1Pattern = 1;
|
||||
GatherMask(sortedGroupTensor, top2ValueInGroupTensor, src1Pattern, false, static_cast<uint32_t>(0), {1, 1, 0, 0},
|
||||
rsvdCnt);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Cast(sortedGroupIndexTensor, sortedGroupTensor, RoundMode::CAST_ROUND, kGroup_);
|
||||
|
||||
sortedGroupQueue_.FreeTensor(sortedGroupTensor);
|
||||
sortedInGroupQueue_.EnQue<float>(sortedInGroupTensor);
|
||||
sigmoidTmpQueue_.FreeTensor(top2ValueInGroupTensor);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashEKFullload<T>::SelectTopKExpertIdx()
|
||||
{
|
||||
LocalTensor<int32_t> expertIdxTensor = expertIdxOutQueue_.AllocTensor<int32_t>();
|
||||
LocalTensor<int32_t> topKGroupIndexTensor = calcTmpBuffer_.Get<int32_t>();
|
||||
LocalTensor<float> sortedInGroupTensor = sortedInGroupQueue_.DeQue<float>();
|
||||
LocalTensor<float> sortedExpertTensor = xInQueue_.AllocTensor<float>();
|
||||
AscendC::MrgSort4Info params;
|
||||
params.elementLengths[0] = k_;
|
||||
params.elementLengths[1] = k_;
|
||||
params.elementLengths[2] = k_;
|
||||
params.elementLengths[3] = k_;
|
||||
params.ifExhaustedSuspension = true;
|
||||
params.validBit = 0b1111;
|
||||
params.repeatTimes = 1;
|
||||
event_t eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
int64_t listOffset1 = topKGroupIndexTensor.GetValue(3) * perGroupExpertCount_ * 2;
|
||||
int64_t listOffset2 = topKGroupIndexTensor.GetValue(2) * perGroupExpertCount_ * 2;
|
||||
int64_t listOffset3 = topKGroupIndexTensor.GetValue(1) * perGroupExpertCount_ * 2;
|
||||
int64_t listOffset4 = topKGroupIndexTensor.GetValue(0) * perGroupExpertCount_ * 2;
|
||||
event_t eventIdSToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
AscendC::MrgSortSrcList<float> srcList;
|
||||
srcList.src1 = sortedInGroupTensor[listOffset1];
|
||||
srcList.src2 = sortedInGroupTensor[listOffset2];
|
||||
srcList.src3 = sortedInGroupTensor[listOffset3];
|
||||
srcList.src4 = sortedInGroupTensor[listOffset4];
|
||||
MrgSort<float>(sortedExpertTensor, srcList, params);
|
||||
PipeBarrier<PIPE_V>();
|
||||
uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数
|
||||
uint8_t src1Pattern = 2; // 内置固定模式
|
||||
GatherMask(expertIdxTensor, sortedExpertTensor.template ReinterpretCast<int32_t>(), src1Pattern, false,
|
||||
static_cast<uint32_t>(0), {1, 1, 0, 0}, rsvdCnt);
|
||||
xInQueue_.FreeTensor(sortedExpertTensor);
|
||||
expertIdxOutQueue_.EnQue(expertIdxTensor);
|
||||
sortedInGroupQueue_.FreeTensor(sortedInGroupTensor);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashEKFullload<T>::SelectTopKExpertScore()
|
||||
{
|
||||
LocalTensor<int32_t> expertIdxTensor = expertIdxOutQueue_.DeQue<int32_t>();
|
||||
LocalTensor<int32_t> expertByteIdxTensor = calcTmpBuffer_.Get<int32_t>();
|
||||
LocalTensor<float> xSigmoidTensor = xSigmoidQueue_.DeQue<float>();
|
||||
LocalTensor<T> yTensor = yOutQueue_.AllocTensor<T>();
|
||||
LocalTensor<float> yOutTensor;
|
||||
if constexpr (!IsSameType<T, float>::value) {
|
||||
yOutTensor = yTensor.template ReinterpretCast<float>()[kAlign_];
|
||||
} else {
|
||||
yOutTensor = yTensor;
|
||||
}
|
||||
Muls(expertByteIdxTensor, expertIdxTensor, static_cast<int32_t>(sizeof(float)), k_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Gather(yOutTensor, xSigmoidTensor, expertByteIdxTensor.template ReinterpretCast<uint32_t>(),
|
||||
static_cast<uint32_t>(0), k_);
|
||||
|
||||
LocalTensor<float> calTensor = calcTmpBuffer_.Get<float>();
|
||||
PipeBarrier<PIPE_V>();
|
||||
ReduceSum(calTensor, yOutTensor, xSigmoidTensor, k_);
|
||||
event_t eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
float sumValue = calTensor.GetValue(0) + eps_;
|
||||
event_t eventIdSToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Duplicate(calTensor, sumValue, k_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Div(yOutTensor, yOutTensor, calTensor, k_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Muls(yOutTensor, yOutTensor, routedScalingFactor_, k_);
|
||||
|
||||
if constexpr (!IsSameType<T, float>::value) {
|
||||
PipeBarrier<PIPE_V>();
|
||||
Cast(yTensor, yOutTensor, RoundMode::CAST_RINT, k_);
|
||||
}
|
||||
|
||||
xSigmoidQueue_.EnQue<float>(xSigmoidTensor);
|
||||
expertIdxOutQueue_.EnQue<int32_t>(expertIdxTensor);
|
||||
yOutQueue_.EnQue(yTensor);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashEKFullload<T>::CopyOut(int64_t row)
|
||||
{
|
||||
LocalTensor<T> yOutTensor = yOutQueue_.DeQue<T>();
|
||||
LocalTensor<int32_t> expertIdxTensor = expertIdxOutQueue_.DeQue<int32_t>();
|
||||
LocalTensor<float> xSigmoidTensor = xSigmoidQueue_.DeQue<float>();
|
||||
DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(k_ * sizeof(T)), 0, 0, 0};
|
||||
DataCopyPad(yGm_[row * k_], yOutTensor, dataCopyParams);
|
||||
dataCopyParams.blockLen = k_ * sizeof(int32_t);
|
||||
DataCopyPad(expertIdxGm_[row * k_], expertIdxTensor, dataCopyParams);
|
||||
xSigmoidQueue_.FreeTensor(xSigmoidTensor);
|
||||
expertIdxOutQueue_.FreeTensor(expertIdxTensor);
|
||||
yOutQueue_.FreeTensor(yOutTensor);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashEKFullload<T>::Init(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx,
|
||||
GM_ADDR out, GM_ADDR workspace,
|
||||
const MoeGatingTopKHashTilingData *tilingData, TPipe *tPipe)
|
||||
{
|
||||
tilingData_ = tilingData;
|
||||
pipe_ = tPipe;
|
||||
blockIdx_ = GetBlockIdx();
|
||||
perCoreRowCount_ = tilingData_->perCoreRowCount;
|
||||
if (blockIdx_ == GetBlockNum() - 1) {
|
||||
curCoreRowCount_ = tilingData_->lastCoreRowCount;
|
||||
} else {
|
||||
curCoreRowCount_ = tilingData_->perCoreRowCount;
|
||||
}
|
||||
expertCount_ = tilingData_->expertCount;
|
||||
addBias_ = tilingData_->addBias == 1;
|
||||
k_ = tilingData_->k;
|
||||
kGroup_ = tilingData_->kGroup;
|
||||
groupCount_ = tilingData_->groupCount;
|
||||
perGroupExpertCount_ = tilingData_->perGroupExpertCount;
|
||||
routedScalingFactor_ = tilingData_->routedScalingFactor;
|
||||
eps_ = tilingData_->eps;
|
||||
|
||||
expertCountAlign_ = Align(expertCount_, sizeof(float));
|
||||
kAlign_ = Align(expertCount_, sizeof(float));
|
||||
|
||||
// init input gm buf
|
||||
xGm_.SetGlobalBuffer((__gm__ T *)x + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_);
|
||||
biasGm_.SetGlobalBuffer((__gm__ T *)bias, expertCount_);
|
||||
|
||||
// init output gm buf
|
||||
yGm_.SetGlobalBuffer((__gm__ T *)y + perCoreRowCount_ * k_ * blockIdx_, k_);
|
||||
expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx + perCoreRowCount_ * k_ * blockIdx_, k_);
|
||||
outGm_.SetGlobalBuffer((__gm__ T *)out + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_);
|
||||
|
||||
// init que
|
||||
pipe_->InitBuffer(xInQueue_, 2, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T)));
|
||||
pipe_->InitBuffer(biasInQueue_, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T)));
|
||||
|
||||
pipe_->InitBuffer(xSigmoidQueue_, 1, AlignBytes(expertCount_, sizeof(float)));
|
||||
pipe_->InitBuffer(xBiasQueue_, 2, AlignBytes(expertCount_, sizeof(float)));
|
||||
|
||||
pipe_->InitBuffer(yOutQueue_, 2, kAlign_ * sizeof(float) * (sizeof(float) / sizeof(T)));
|
||||
pipe_->InitBuffer(expertIdxOutQueue_, 2, AlignBytes(k_, sizeof(int32_t)));
|
||||
pipe_->InitBuffer(outOutQueue_, 2, AlignBytes(expertCount_, sizeof(float)));
|
||||
|
||||
pipe_->InitBuffer(sigmoidTmpQueue_, 2, AlignBytes(expertCount_, sizeof(float)));
|
||||
pipe_->InitBuffer(sortedInGroupQueue_, 2, AlignBytes(expertCount_, sizeof(float)) * 2);
|
||||
pipe_->InitBuffer(sortedGroupQueue_, 2,
|
||||
(groupCount_ + ONE_REPEAT_SORT_NUM - 1) / ONE_REPEAT_SORT_NUM * ONE_REPEAT_SORT_NUM *
|
||||
sizeof(float) * 2);
|
||||
|
||||
pipe_->InitBuffer(calcTmpBuffer_, tilingData_->calTmpBufUbSize);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashEKFullload<T>::Process()
|
||||
{
|
||||
CopyInBias();
|
||||
for (int64_t row = 0; row < curCoreRowCount_; row++) {
|
||||
CopyInX(row);
|
||||
ComputeX();
|
||||
SortInGroup();
|
||||
SelectTopKGroupIndex();
|
||||
SelectTopKExpertIdx();
|
||||
SelectTopKExpertScore();
|
||||
CopyOut(row);
|
||||
}
|
||||
}
|
||||
} // namespace MoeGatingTopKHash
|
||||
#endif // MOE_GATING_TOP_K_E_K_FULLLOAD_H
|
||||
@@ -0,0 +1,664 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file moe_gating_top_k_hash_generalized.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef MOE_GATING_TOP_K_E_K_GENERALIZED_H
|
||||
#define MOE_GATING_TOP_K_E_K_GENERALIZED_H
|
||||
#include "kernel_operator.h"
|
||||
#include "common.h"
|
||||
#include "kernel_utils.h"
|
||||
namespace MoeGatingTopKHash {
|
||||
using namespace AscendC;
|
||||
|
||||
template <typename T>
|
||||
class MoeGatingTopKHashGenerlized {
|
||||
public:
|
||||
__aicore__ inline MoeGatingTopKHashGenerlized(){};
|
||||
__aicore__ inline void Init(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, GM_ADDR workspace,
|
||||
const MoeGatingTopKHashTilingData *tilingData, TPipe *tPipe);
|
||||
__aicore__ inline void Process();
|
||||
|
||||
private:
|
||||
__aicore__ inline void CopyInBiasAndInitExpertId();
|
||||
__aicore__ inline void CopyInX(int64_t progress);
|
||||
__aicore__ inline void ComputeX();
|
||||
__aicore__ inline void CopuOutXNorm(int64_t row);
|
||||
__aicore__ inline void SortInGroup();
|
||||
__aicore__ inline void SelectTopKGroupIndex();
|
||||
__aicore__ inline void SelectTopKExpertIdx();
|
||||
__aicore__ inline void SelectTopKExpertScore();
|
||||
__aicore__ inline void CumputeActualTopKExpertId();
|
||||
__aicore__ inline void CopyOut(int64_t row);
|
||||
|
||||
private:
|
||||
TPipe *pipe_;
|
||||
TQue<QuePosition::VECIN, 1> xInQueue_;
|
||||
TQue<QuePosition::VECOUT, 1> yOutQueue_;
|
||||
TQue<QuePosition::VECOUT, 1> expertIdxOutQueue_;
|
||||
TQue<QuePosition::VECOUT, 1> outOutQueue_;
|
||||
|
||||
TBuf<TPosition::VECCALC> biasBuf_; // 存放输入bias
|
||||
TBuf<TPosition::VECCALC> expertIdBuf_; // 专家编号
|
||||
TBuf<TPosition::VECCALC> xNormWithBiasBuf_; // 存放加了bias之后的值
|
||||
TBuf<TPosition::VECCALC> xNormBuf_; // 存放计算sigmoid或softmax的值
|
||||
TBuf<TPosition::VECCALC> sortedInGroupBuf_; // 存放组内排序后的结果
|
||||
TBuf<TPosition::VECCALC> topKExpertIdBuf_;
|
||||
TBuf<TPosition::VECCALC> sortedGroupIndexBuf_;
|
||||
TBuf<TPosition::VECCALC> calcTmpBuf_;
|
||||
|
||||
GlobalTensor<T> xGm_;
|
||||
GlobalTensor<T> biasGm_;
|
||||
GlobalTensor<T> yGm_;
|
||||
GlobalTensor<int32_t> expertIdxGm_;
|
||||
GlobalTensor<float> outGm_;
|
||||
|
||||
int64_t blockIdx_ = 0;
|
||||
int64_t perCoreRowCount_ = 0;
|
||||
int64_t curCoreRowCount_ = 0;
|
||||
int64_t expertCount_ = 0;
|
||||
bool addBias_ = false;
|
||||
int64_t k_ = 0;
|
||||
int64_t kGroup_ = 0;
|
||||
int64_t groupCount_ = 0;
|
||||
int64_t groupCountAlign_ = 0;
|
||||
int64_t perGroupExpertCount_ = 0;
|
||||
int64_t perGroupExpertCountAlign_ = 0;
|
||||
int64_t groupSelectMode_ = 0;
|
||||
int64_t renorm_ = 0;
|
||||
int64_t normType_ = 0;
|
||||
int64_t outFlag_ = 0;
|
||||
|
||||
int64_t expertCountAlign_ = 0;
|
||||
int64_t kAlign_ = 0;
|
||||
bool isAlign_ = false;
|
||||
|
||||
const MoeGatingTopKHashTilingData *tilingData_;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashGenerlized<T>::CopyInBiasAndInitExpertId()
|
||||
{
|
||||
LocalTensor<float> biasTensor = biasBuf_.Get<float>();
|
||||
LocalTensor<int32_t> expertIdTensor = expertIdBuf_.Get<int32_t>();
|
||||
DataCopyExtParams dataCopyParams;
|
||||
dataCopyParams.blockCount = groupCount_;
|
||||
dataCopyParams.blockLen = perGroupExpertCount_ * sizeof(T);
|
||||
dataCopyParams.srcStride = 0;
|
||||
dataCopyParams.dstStride = (perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(T) / BLOCK_BYTES;
|
||||
|
||||
DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast<T>(0)};
|
||||
if (addBias_) {
|
||||
if constexpr (IsSameType<T, float>::value) {
|
||||
DataCopyPad(biasTensor, biasGm_, dataCopyParams, dataCopyPadParams);
|
||||
event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
|
||||
SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
|
||||
WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
|
||||
} else {
|
||||
DataCopyPad(biasTensor[expertCountAlign_].ReinterpretCast<T>(), biasGm_, dataCopyParams, dataCopyPadParams);
|
||||
event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
|
||||
SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
|
||||
WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
|
||||
Cast(biasTensor, biasTensor[expertCountAlign_].ReinterpretCast<T>(), RoundMode::CAST_NONE,
|
||||
expertCountAlign_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
|
||||
if (!isAlign_) {
|
||||
int64_t duplicateNum = perGroupExpertCount_ % ONE_REPEAT_SORT_NUM;
|
||||
int duplicateIndex = perGroupExpertCount_ - duplicateNum;
|
||||
if (duplicateNum > 0) {
|
||||
uint64_t mask0 = UINT64_MAX;
|
||||
mask0 = mask0 << duplicateNum;
|
||||
mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
|
||||
uint64_t mask[2] = {mask0, 0};
|
||||
Duplicate(biasTensor.ReinterpretCast<int32_t>()[duplicateIndex], FLOAT32_NEG_INF, mask, groupCount_, 1,
|
||||
perGroupExpertCountAlign_ * sizeof(float) / BLOCK_BYTES);
|
||||
}
|
||||
}
|
||||
}
|
||||
ArithProgression(expertIdTensor, static_cast<int32_t>(0), static_cast<int32_t>(1), expertCountAlign_);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashGenerlized<T>::CopyInX(int64_t row)
|
||||
{
|
||||
LocalTensor<float> xInLocalTensor = xInQueue_.AllocTensor<float>();
|
||||
DataCopyExtParams dataCopyParams;
|
||||
dataCopyParams.blockCount = groupCount_;
|
||||
dataCopyParams.blockLen = perGroupExpertCount_ * sizeof(T);
|
||||
dataCopyParams.srcStride = 0;
|
||||
dataCopyParams.dstStride = (perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(T) / BLOCK_BYTES;
|
||||
|
||||
DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast<T>(0)};
|
||||
if constexpr (IsSameType<T, float>::value) {
|
||||
DataCopyPad(xInLocalTensor, xGm_[row * expertCount_], dataCopyParams, dataCopyPadParams);
|
||||
} else {
|
||||
DataCopyPad(xInLocalTensor[expertCountAlign_].ReinterpretCast<T>(), xGm_[row * expertCount_], dataCopyParams,
|
||||
dataCopyPadParams);
|
||||
}
|
||||
xInQueue_.EnQue(xInLocalTensor);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashGenerlized<T>::ComputeX()
|
||||
{
|
||||
LocalTensor<float> xNormTensor = xNormBuf_.Get<float>();
|
||||
LocalTensor<float> xInLocalTensor = xInQueue_.DeQue<float>();
|
||||
LocalTensor<float> xNormWithBiasTensor = xNormWithBiasBuf_.Get<float>();
|
||||
LocalTensor<float> biasTensor = biasBuf_.Get<float>();
|
||||
|
||||
if constexpr (!IsSameType<T, float>::value) {
|
||||
Cast(xInLocalTensor, xInLocalTensor[expertCountAlign_].ReinterpretCast<T>(), RoundMode::CAST_NONE,
|
||||
expertCountAlign_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
|
||||
int64_t duplicateNum = perGroupExpertCount_ % ONE_REPEAT_SORT_NUM;
|
||||
int duplicateIndex = perGroupExpertCount_ - duplicateNum;
|
||||
if (!isAlign_ && duplicateNum > 0) {
|
||||
uint64_t mask0 = UINT64_MAX;
|
||||
mask0 = mask0 << duplicateNum;
|
||||
mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
|
||||
uint64_t mask[2] = {mask0, 0};
|
||||
Duplicate(xInLocalTensor.ReinterpretCast<int32_t>()[duplicateIndex], FLOAT32_NEG_INF, mask, groupCount_, 1,
|
||||
(perGroupExpertCountAlign_ * sizeof(float)) / BLOCK_BYTES);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
if (tilingData_->normType == 1) { // sigmoid
|
||||
LocalTensor<uint8_t> calcNormTmpTensor = calcTmpBuf_.Get<uint8_t>();
|
||||
Sigmoid(xNormTensor, xInLocalTensor, calcNormTmpTensor, expertCountAlign_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
} else { // softmax
|
||||
LocalTensor<float> reduceValueTensor = calcTmpBuf_.Get<float>();
|
||||
LocalTensor<float> calcTmp = calcTmpBuf_.Get<float>()[BLOCK_BYTES];
|
||||
ReduceMax(reduceValueTensor, xInLocalTensor, calcTmp, expertCountAlign_);
|
||||
event_t eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
float maxValue = reduceValueTensor.GetValue(0);
|
||||
event_t eventIdSToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Adds(xNormTensor, xInLocalTensor, -maxValue, expertCountAlign_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Exp(xNormTensor, xNormTensor, expertCountAlign_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
ReduceSum(reduceValueTensor, xNormTensor, calcTmp, expertCountAlign_);
|
||||
eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
float sumValue = reduceValueTensor.GetValue(0);
|
||||
eventIdSToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Muls(xNormTensor, xNormTensor, 1.0f / sumValue, expertCountAlign_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
if (addBias_) {
|
||||
Add(xNormWithBiasTensor, xNormTensor, biasTensor, expertCountAlign_);
|
||||
} else {
|
||||
DataCopy(xNormWithBiasTensor, xNormTensor, expertCountAlign_);
|
||||
}
|
||||
|
||||
if (!isAlign_ && duplicateNum > 0) {
|
||||
uint64_t mask0 = UINT64_MAX;
|
||||
mask0 = mask0 << duplicateNum;
|
||||
mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
|
||||
uint64_t mask[2] = {mask0, 0};
|
||||
PipeBarrier<PIPE_V>();
|
||||
Duplicate(xNormWithBiasTensor.ReinterpretCast<int32_t>()[duplicateIndex],
|
||||
FLOAT32_NEG_INF, // MIN_FP32,
|
||||
mask, groupCount_, 1, perGroupExpertCountAlign_ * sizeof(float) / BLOCK_BYTES);
|
||||
}
|
||||
xInQueue_.FreeTensor(xInLocalTensor);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashGenerlized<T>::CopuOutXNorm(int64_t row)
|
||||
{
|
||||
LocalTensor<float> outOutTensor = outOutQueue_.AllocTensor<float>();
|
||||
LocalTensor<float> xNormTensor = xNormBuf_.Get<float>();
|
||||
DataCopy(outOutTensor, xNormTensor, expertCountAlign_);
|
||||
outOutQueue_.EnQue<float>(outOutTensor);
|
||||
outOutTensor = outOutQueue_.DeQue<float>();
|
||||
DataCopyExtParams dataCopyParams{
|
||||
static_cast<uint16_t>(groupCount_), static_cast<uint32_t>(perGroupExpertCount_ * sizeof(float)),
|
||||
static_cast<uint32_t>((perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(float) / BLOCK_BYTES), 0, 0};
|
||||
DataCopyPad(outGm_[row * expertCount_], outOutTensor, dataCopyParams);
|
||||
outOutQueue_.FreeTensor(outOutTensor);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashGenerlized<T>::SortInGroup()
|
||||
{
|
||||
LocalTensor<float> xNormWithBiasTensor = xNormWithBiasBuf_.Get<float>();
|
||||
LocalTensor<uint32_t> expertIdTensor = expertIdBuf_.Get<uint32_t>();
|
||||
LocalTensor<float> sortedInGroupTensor = sortedInGroupBuf_.Get<float>();
|
||||
LocalTensor<float> tmpLocal = calcTmpBuf_.Get<float>();
|
||||
if (perGroupExpertCountAlign_ == ONE_REPEAT_SORT_NUM) {
|
||||
PipeBarrier<PIPE_V>();
|
||||
Sort32(sortedInGroupTensor, xNormWithBiasTensor, expertIdTensor, groupCount_);
|
||||
} else {
|
||||
for (int64_t group = 0; group < groupCount_; group++) {
|
||||
PipeBarrier<PIPE_V>();
|
||||
Sort<float, true>(sortedInGroupTensor[group * perGroupExpertCountAlign_ * CONSTANT_TWO],
|
||||
xNormWithBiasTensor[group * perGroupExpertCountAlign_],
|
||||
expertIdTensor[group * perGroupExpertCountAlign_], tmpLocal,
|
||||
perGroupExpertCountAlign_ / ONE_REPEAT_SORT_NUM);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashGenerlized<T>::SelectTopKGroupIndex()
|
||||
{
|
||||
LocalTensor<float> sortedInGroupTensor = sortedInGroupBuf_.Get<float>();
|
||||
LocalTensor<float> valueSelectedFromGroupTensor = calcTmpBuf_.GetWithOffset<float>(groupCountAlign_ * 2, 0);
|
||||
LocalTensor<uint32_t> maskTensor =
|
||||
calcTmpBuf_.GetWithOffset<uint32_t>(groupCountAlign_, groupCountAlign_ * 2 * sizeof(float));
|
||||
LocalTensor<float> topValueInGroupTensor =
|
||||
calcTmpBuf_.GetWithOffset<float>(groupCountAlign_, groupCountAlign_ * 3 * sizeof(float));
|
||||
LocalTensor<uint32_t> groupIndex =
|
||||
calcTmpBuf_.GetWithOffset<uint32_t>(groupCountAlign_, groupCountAlign_ * 4 * sizeof(float));
|
||||
LocalTensor<float> sortedTopValue =
|
||||
calcTmpBuf_.GetWithOffset<float>(groupCountAlign_ * 2, groupCountAlign_ * 5 * sizeof(float));
|
||||
LocalTensor<float> sortTmp =
|
||||
calcTmpBuf_.GetWithOffset<float>(groupCountAlign_ * 2, groupCountAlign_ * 7 * sizeof(float));
|
||||
|
||||
event_t eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
|
||||
uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数
|
||||
PipeBarrier<PIPE_V>();
|
||||
if (tilingData_->groupSelectMode == 1) { // top2 sum
|
||||
// 提取每组组前两个元素
|
||||
maskTensor.SetValue(0, static_cast<uint32_t>(5)); // b0101
|
||||
maskTensor.SetValue(1, static_cast<uint32_t>(0));
|
||||
event_t eventIdSToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
|
||||
GatherMaskParams gatherMaskParams;
|
||||
gatherMaskParams.repeatTimes = groupCount_;
|
||||
gatherMaskParams.src0BlockStride = 1;
|
||||
gatherMaskParams.src0RepeatStride =
|
||||
Ceil(perGroupExpertCountAlign_ * (sizeof(float) + sizeof(uint32_t)), BLOCK_BYTES);
|
||||
gatherMaskParams.src1RepeatStride = 0;
|
||||
GatherMask(valueSelectedFromGroupTensor, sortedInGroupTensor, maskTensor, true,
|
||||
static_cast<uint32_t>(ONE_REPEAT_SORT_NUM * CONSTANT_TWO), gatherMaskParams, rsvdCnt);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
// 计算每个组前两个数的和
|
||||
PairReduceSum(topValueInGroupTensor, valueSelectedFromGroupTensor,
|
||||
Ceil(groupCount_ * sizeof(float) * 2, REPEAT_BYTES), REPEAT_BYTES / sizeof(float), 1, 1,
|
||||
CONSTANT_EIGHT); // 计算每个组内最大的两个数之和
|
||||
} else {
|
||||
maskTensor.SetValue(0, static_cast<uint32_t>(1)); // b0101
|
||||
maskTensor.SetValue(1, static_cast<uint32_t>(0));
|
||||
|
||||
event_t eventIdSToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数
|
||||
GatherMaskParams gatherMaskParams;
|
||||
gatherMaskParams.repeatTimes = groupCount_;
|
||||
gatherMaskParams.src0BlockStride = 1;
|
||||
gatherMaskParams.src0RepeatStride = Ceil(perGroupExpertCountAlign_ * (sizeof(float) + sizeof(uint32_t)), 32);
|
||||
gatherMaskParams.src1RepeatStride = 0;
|
||||
GatherMask(topValueInGroupTensor, sortedInGroupTensor, maskTensor, true,
|
||||
static_cast<uint32_t>(ONE_REPEAT_SORT_NUM * CONSTANT_TWO), gatherMaskParams, rsvdCnt);
|
||||
}
|
||||
|
||||
PipeBarrier<PIPE_V>();
|
||||
// 生成组索引
|
||||
ArithProgression(groupIndex.ReinterpretCast<int32_t>(), static_cast<int32_t>(0), static_cast<int32_t>(1),
|
||||
groupCount_); // 生成组索引
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
int64_t duplicateNum = groupCount_ % ONE_REPEAT_SORT_NUM;
|
||||
int duplicateIndex = groupCount_ - duplicateNum;
|
||||
if (duplicateNum > 0) {
|
||||
uint64_t mask0 = UINT64_MAX;
|
||||
mask0 = mask0 << duplicateNum;
|
||||
mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
|
||||
uint64_t mask[2] = {mask0, 0};
|
||||
Duplicate(topValueInGroupTensor.ReinterpretCast<int32_t>()[duplicateIndex], FLOAT32_NEG_INF, mask, 1, 1,
|
||||
REPEAT_BLOCKS);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
// 排序
|
||||
Sort<float, true>(sortedTopValue, topValueInGroupTensor, groupIndex, sortTmp, Ceil(groupCount_, 32));
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
// 提取组序号
|
||||
uint8_t src1Pattern = 2; // 内置固定模式
|
||||
GatherMask(groupIndex, sortedTopValue.template ReinterpretCast<uint32_t>(), src1Pattern, false,
|
||||
static_cast<uint32_t>(0),
|
||||
{1, static_cast<uint8_t>(Ceil(kGroup_ * sizeof(float) * CONSTANT_TWO, 256)), REPEAT_BLOCKS, 0}, rsvdCnt);
|
||||
PipeBarrier<PIPE_V>();
|
||||
duplicateNum = kGroup_ % ONE_REPEAT_SORT_NUM;
|
||||
if (duplicateNum > 0) {
|
||||
duplicateIndex = kGroup_ - duplicateNum;
|
||||
uint64_t mask0 = UINT64_MAX;
|
||||
mask0 = mask0 << duplicateNum;
|
||||
mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
|
||||
uint64_t mask[2] = {mask0, 0};
|
||||
PipeBarrier<PIPE_V>();
|
||||
Duplicate(groupIndex.ReinterpretCast<int32_t>()[duplicateIndex], FLOAT32_NEG_INF, mask, 1, 1, REPEAT_BLOCKS);
|
||||
}
|
||||
|
||||
// 将筛选出来的组序号降序排列
|
||||
LocalTensor<float> sortedGroupIndex = sortedGroupIndexBuf_.Get<float>();
|
||||
PipeBarrier<PIPE_V>();
|
||||
Sort<float, true>(sortedGroupIndex, groupIndex.ReinterpretCast<float>(), groupIndex, sortTmp, Ceil(kGroup_, 32));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashGenerlized<T>::SelectTopKExpertIdx()
|
||||
{
|
||||
LocalTensor<float> sortedInGroupTensor = sortedInGroupBuf_.Get<float>();
|
||||
LocalTensor<int32_t> sortedGroupIndex = sortedGroupIndexBuf_.Get<int32_t>();
|
||||
LocalTensor<int32_t> topKExpertId = topKExpertIdBuf_.Get<int32_t>();
|
||||
LocalTensor<float> mrgSort0Tensor = calcTmpBuf_.Get<float>();
|
||||
|
||||
uint32_t offset[CONSTANT_FOUR] = {0, 0, 0, 0};
|
||||
uint16_t lenArr[CONSTANT_FOUR] = {
|
||||
static_cast<uint16_t>(perGroupExpertCount_), static_cast<uint16_t>(perGroupExpertCount_),
|
||||
static_cast<uint16_t>(perGroupExpertCount_), static_cast<uint16_t>(perGroupExpertCount_)};
|
||||
MrgSort4Info params{lenArr, false, 0b1111, 1};
|
||||
MrgSortSrcList<float> srcList;
|
||||
|
||||
event_t eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
|
||||
for (int32_t i = kGroup_ - 1; i >= 0; i -= CONSTANT_FOUR) {
|
||||
int64_t mrgLen = Min(i + 1, CONSTANT_FOUR);
|
||||
if (mrgLen > 1) {
|
||||
if (mrgLen == MERGE_LIST_FOUR) {
|
||||
offset[0] = sortedGroupIndex.GetValue(i * 2) * perGroupExpertCountAlign_ * 2;
|
||||
offset[1] = sortedGroupIndex.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2;
|
||||
offset[2] = sortedGroupIndex.GetValue((i - 2) * 2) * perGroupExpertCountAlign_ * 2;
|
||||
offset[3] = sortedGroupIndex.GetValue((i - 3) * 2) * perGroupExpertCountAlign_ * 2;
|
||||
} else if (mrgLen == MERGE_LIST_THREE) {
|
||||
offset[0] = sortedGroupIndex.GetValue(i * 2) * perGroupExpertCountAlign_ * 2;
|
||||
offset[1] = sortedGroupIndex.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2;
|
||||
offset[2] = sortedGroupIndex.GetValue((i - 2) * 2) * perGroupExpertCountAlign_ * 2;
|
||||
offset[3] = 0;
|
||||
params.elementLengths[3] = 0;
|
||||
params.validBit = 0b111;
|
||||
} else {
|
||||
offset[0] = sortedGroupIndex.GetValue(i * 2) * perGroupExpertCountAlign_ * 2;
|
||||
offset[1] = sortedGroupIndex.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2;
|
||||
offset[2] = 0;
|
||||
offset[3] = 0;
|
||||
params.elementLengths[2] = 0;
|
||||
params.elementLengths[3] = 0;
|
||||
params.validBit = 0b11;
|
||||
}
|
||||
|
||||
srcList.src1 = sortedInGroupTensor[offset[0]];
|
||||
srcList.src2 = sortedInGroupTensor[offset[1]];
|
||||
srcList.src3 = sortedInGroupTensor[offset[2]];
|
||||
srcList.src4 = sortedInGroupTensor[offset[3]];
|
||||
|
||||
PipeBarrier<PIPE_V>();
|
||||
MrgSort(mrgSort0Tensor[(kGroup_ - 1 - i) * perGroupExpertCountAlign_ * 2], srcList, params);
|
||||
} else {
|
||||
offset[0] = sortedGroupIndex.GetValue(i * 2) * perGroupExpertCountAlign_ * 2;
|
||||
PipeBarrier<PIPE_V>();
|
||||
DataCopy(mrgSort0Tensor[(kGroup_ - 1 - i) * perGroupExpertCountAlign_ * 2], sortedInGroupTensor[offset[0]],
|
||||
perGroupExpertCountAlign_ * 2);
|
||||
}
|
||||
}
|
||||
int32_t baseLoop = 4;
|
||||
LocalTensor<float> srcTensor = mrgSort0Tensor;
|
||||
LocalTensor<float> dstTensor = mrgSort0Tensor;
|
||||
for (int i = 0; i < tilingData_->vmsCount; i++) {
|
||||
if (i % 2 == 0) {
|
||||
srcTensor = mrgSort0Tensor;
|
||||
dstTensor = sortedInGroupTensor;
|
||||
} else {
|
||||
srcTensor = sortedInGroupTensor;
|
||||
dstTensor = mrgSort0Tensor;
|
||||
}
|
||||
|
||||
int32_t nextBaseRow = baseLoop * MERGE_LIST_FOUR;
|
||||
int32_t quotient = kGroup_ / nextBaseRow;
|
||||
int32_t remainder = kGroup_ - quotient * nextBaseRow;
|
||||
if (quotient > 0) {
|
||||
MrgSort4Info params;
|
||||
MrgSortSrcList<float> srcList;
|
||||
params.ifExhaustedSuspension = false;
|
||||
params.elementLengths[0] = perGroupExpertCount_ * baseLoop;
|
||||
params.elementLengths[1] = perGroupExpertCount_ * baseLoop;
|
||||
params.elementLengths[2] = perGroupExpertCount_ * baseLoop;
|
||||
params.elementLengths[3] = perGroupExpertCount_ * baseLoop;
|
||||
params.validBit = 0b1111;
|
||||
params.repeatTimes = 1;
|
||||
for (int j = 0; j < quotient; j++) {
|
||||
srcList.src1 = srcTensor[perGroupExpertCountAlign_ * baseLoop * 8 * j];
|
||||
srcList.src2 = srcTensor[perGroupExpertCountAlign_ * baseLoop * (8 * j + 2)];
|
||||
srcList.src3 = srcTensor[perGroupExpertCountAlign_ * baseLoop * (8 * j + 4)];
|
||||
srcList.src4 = srcTensor[perGroupExpertCountAlign_ * baseLoop * (8 * j + 6)];
|
||||
PipeBarrier<PIPE_V>();
|
||||
MrgSort(dstTensor[perGroupExpertCountAlign_ * baseLoop * 8 * j], srcList, params);
|
||||
}
|
||||
}
|
||||
|
||||
if (remainder > 0) {
|
||||
int32_t baseOffset = quotient * nextBaseRow * perGroupExpertCountAlign_ * 2;
|
||||
int32_t mrgLen = CeilDiv(remainder, baseLoop);
|
||||
int32_t tailRow = remainder - (mrgLen - 1) * baseLoop;
|
||||
if (mrgLen > 1) {
|
||||
MrgSort4Info params;
|
||||
MrgSortSrcList<float> srcList;
|
||||
params.repeatTimes = 1;
|
||||
params.ifExhaustedSuspension = false;
|
||||
params.elementLengths[0] = perGroupExpertCount_ * baseLoop;
|
||||
params.elementLengths[1] = perGroupExpertCount_ * baseLoop;
|
||||
params.elementLengths[2] = perGroupExpertCount_ * baseLoop;
|
||||
params.elementLengths[3] = perGroupExpertCount_ * baseLoop;
|
||||
srcList.src1 = srcTensor[baseOffset];
|
||||
srcList.src2 = srcTensor[baseOffset + perGroupExpertCountAlign_ * baseLoop * 2];
|
||||
if (mrgLen == MERGE_LIST_FOUR) {
|
||||
srcList.src3 = srcTensor[baseOffset + perGroupExpertCountAlign_ * baseLoop * 2 * 2];
|
||||
srcList.src4 = srcTensor[baseOffset + perGroupExpertCountAlign_ * baseLoop * 2 * 3];
|
||||
params.elementLengths[3] = perGroupExpertCount_ * tailRow;
|
||||
params.validBit = 0b1111;
|
||||
} else if (mrgLen == MERGE_LIST_THREE) {
|
||||
srcList.src3 = srcTensor[baseOffset + perGroupExpertCountAlign_ * baseLoop * 2 * 2];
|
||||
params.elementLengths[2] = perGroupExpertCount_ * tailRow;
|
||||
params.elementLengths[3] = 0;
|
||||
params.validBit = 0b111;
|
||||
} else {
|
||||
params.elementLengths[1] = perGroupExpertCount_ * tailRow;
|
||||
params.elementLengths[2] = 0;
|
||||
params.elementLengths[3] = 0;
|
||||
params.validBit = 0b11;
|
||||
}
|
||||
PipeBarrier<PIPE_V>();
|
||||
MrgSort(dstTensor[baseOffset], srcList, params);
|
||||
} else {
|
||||
PipeBarrier<PIPE_V>();
|
||||
DataCopy(dstTensor[baseOffset], srcTensor[baseOffset], tailRow * perGroupExpertCountAlign_ * 2);
|
||||
}
|
||||
}
|
||||
baseLoop = nextBaseRow;
|
||||
}
|
||||
|
||||
GatherMaskParams gatherMaskParams;
|
||||
gatherMaskParams.repeatTimes = Ceil(k_ * sizeof(float) * 2, REPEAT_BYTES);
|
||||
gatherMaskParams.src0BlockStride = 1;
|
||||
gatherMaskParams.src0RepeatStride = REPEAT_BLOCKS;
|
||||
gatherMaskParams.src1RepeatStride = 0;
|
||||
|
||||
uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数
|
||||
uint8_t src1Pattern = 2; // 内置固定模式
|
||||
PipeBarrier<PIPE_V>();
|
||||
GatherMask(topKExpertId, dstTensor.template ReinterpretCast<int32_t>(), src1Pattern, false,
|
||||
static_cast<uint32_t>(0), gatherMaskParams, rsvdCnt);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashGenerlized<T>::SelectTopKExpertScore()
|
||||
{
|
||||
LocalTensor<float> xNormTensor = xNormBuf_.Get<float>();
|
||||
LocalTensor<float> yOutTensor = yOutQueue_.AllocTensor<float>();
|
||||
LocalTensor<int32_t> topKExpertId = topKExpertIdBuf_.Get<int32_t>();
|
||||
LocalTensor<int32_t> topKExpertIdWithByte = calcTmpBuf_.Get<int32_t>();
|
||||
PipeBarrier<PIPE_V>();
|
||||
Muls(topKExpertIdWithByte, topKExpertId, static_cast<int32_t>(sizeof(float)), k_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Gather(yOutTensor, xNormTensor, topKExpertIdWithByte.template ReinterpretCast<uint32_t>(), static_cast<uint32_t>(0),
|
||||
k_);
|
||||
|
||||
if (tilingData_->normType == 1) {
|
||||
LocalTensor<float> maxValueTensor = calcTmpBuf_.Get<float>();
|
||||
LocalTensor<float> tmpTensor = calcTmpBuf_.Get<float>()[32];
|
||||
PipeBarrier<PIPE_V>();
|
||||
ReduceSum(maxValueTensor, yOutTensor, tmpTensor, k_);
|
||||
event_t eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
float sumValue = maxValueTensor.GetValue(0) + tilingData_->eps;
|
||||
event_t eventIdSToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Duplicate(tmpTensor, sumValue, k_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Div(yOutTensor, yOutTensor, tmpTensor, k_);
|
||||
}
|
||||
PipeBarrier<PIPE_V>();
|
||||
Muls(yOutTensor, yOutTensor, tilingData_->routedScalingFactor, k_);
|
||||
|
||||
if constexpr (!IsSameType<T, float>::value) {
|
||||
PipeBarrier<PIPE_V>();
|
||||
Cast(yOutTensor.ReinterpretCast<T>(), yOutTensor, RoundMode::CAST_RINT, k_);
|
||||
}
|
||||
|
||||
yOutQueue_.EnQue<float>(yOutTensor);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashGenerlized<T>::CumputeActualTopKExpertId()
|
||||
{
|
||||
LocalTensor<int32_t> expertIdxOut = expertIdxOutQueue_.AllocTensor<int32_t>();
|
||||
LocalTensor<int32_t> topKExpertId = topKExpertIdBuf_.Get<int32_t>();
|
||||
LocalTensor<float> topKExpertIdFp32 = calcTmpBuf_.Get<float>();
|
||||
|
||||
PipeBarrier<PIPE_V>();
|
||||
Cast(topKExpertIdFp32, topKExpertId, RoundMode::CAST_ROUND, k_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Muls(topKExpertIdFp32, topKExpertIdFp32, 1.0f / (float)perGroupExpertCountAlign_, k_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Cast(expertIdxOut, topKExpertIdFp32, RoundMode::CAST_TRUNC, k_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Muls(expertIdxOut, expertIdxOut, static_cast<int32_t>(perGroupExpertCountAlign_ - perGroupExpertCount_), k_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Sub(expertIdxOut, topKExpertId, expertIdxOut, k_);
|
||||
expertIdxOutQueue_.EnQue<int32_t>(expertIdxOut);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashGenerlized<T>::CopyOut(int64_t row)
|
||||
{
|
||||
LocalTensor<T> yOutTensor = yOutQueue_.DeQue<T>();
|
||||
LocalTensor<int32_t> expertIdxOut = expertIdxOutQueue_.DeQue<int32_t>();
|
||||
DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(k_ * sizeof(T)), 0, 0, 0};
|
||||
DataCopyPad(yGm_[row * k_], yOutTensor, dataCopyParams);
|
||||
dataCopyParams.blockLen = k_ * sizeof(int32_t);
|
||||
DataCopyPad(expertIdxGm_[row * k_], expertIdxOut, dataCopyParams);
|
||||
yOutQueue_.FreeTensor(yOutTensor);
|
||||
expertIdxOutQueue_.FreeTensor(expertIdxOut);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashGenerlized<T>::Init(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx,
|
||||
GM_ADDR out, GM_ADDR workspace,
|
||||
const MoeGatingTopKHashTilingData *tilingData, TPipe *tPipe)
|
||||
{
|
||||
tilingData_ = tilingData;
|
||||
pipe_ = tPipe;
|
||||
blockIdx_ = GetBlockIdx();
|
||||
perCoreRowCount_ = tilingData_->perCoreRowCount;
|
||||
if (blockIdx_ == GetBlockNum() - 1) {
|
||||
curCoreRowCount_ = tilingData_->lastCoreRowCount;
|
||||
} else {
|
||||
curCoreRowCount_ = tilingData_->perCoreRowCount;
|
||||
}
|
||||
expertCount_ = tilingData_->expertCount;
|
||||
addBias_ = tilingData_->addBias == 1;
|
||||
k_ = tilingData_->k;
|
||||
kGroup_ = tilingData_->kGroup;
|
||||
groupCount_ = tilingData_->groupCount;
|
||||
groupCountAlign_ = Ceil(groupCount_, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
|
||||
perGroupExpertCount_ = tilingData_->perGroupExpertCount;
|
||||
perGroupExpertCountAlign_ = tilingData_->perGroupExpertCountAlign;
|
||||
|
||||
expertCountAlign_ = Align(perGroupExpertCountAlign_ * groupCount_, sizeof(float));
|
||||
kAlign_ = Align(k_, sizeof(float));
|
||||
|
||||
isAlign_ = perGroupExpertCount_ == perGroupExpertCountAlign_;
|
||||
|
||||
// init input gm buf
|
||||
xGm_.SetGlobalBuffer((__gm__ T *)x + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_);
|
||||
biasGm_.SetGlobalBuffer((__gm__ T *)bias, expertCount_);
|
||||
|
||||
// init output gm buf
|
||||
yGm_.SetGlobalBuffer((__gm__ T *)y + perCoreRowCount_ * k_ * blockIdx_, k_);
|
||||
expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx + perCoreRowCount_ * k_ * blockIdx_, k_);
|
||||
outGm_.SetGlobalBuffer((__gm__ float *)out + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_);
|
||||
|
||||
// init que
|
||||
pipe_->InitBuffer(xInQueue_, 1, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T)));
|
||||
pipe_->InitBuffer(yOutQueue_, 1, kAlign_ * sizeof(float));
|
||||
pipe_->InitBuffer(expertIdxOutQueue_, 1, kAlign_ * sizeof(int32_t));
|
||||
pipe_->InitBuffer(outOutQueue_, 1, expertCountAlign_ * sizeof(float));
|
||||
|
||||
pipe_->InitBuffer(biasBuf_, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T)));
|
||||
pipe_->InitBuffer(expertIdBuf_, expertCountAlign_ * sizeof(int32_t));
|
||||
|
||||
pipe_->InitBuffer(xNormBuf_, expertCountAlign_ * sizeof(float));
|
||||
|
||||
pipe_->InitBuffer(xNormWithBiasBuf_, expertCountAlign_ * sizeof(float));
|
||||
pipe_->InitBuffer(sortedInGroupBuf_, expertCountAlign_ * (sizeof(float) + sizeof(uint32_t)));
|
||||
|
||||
pipe_->InitBuffer(sortedGroupIndexBuf_, groupCountAlign_ * sizeof(float) * CONSTANT_TWO);
|
||||
pipe_->InitBuffer(topKExpertIdBuf_, kAlign_ * sizeof(int32_t));
|
||||
pipe_->InitBuffer(calcTmpBuf_, expertCountAlign_ * sizeof(float) * 10);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void MoeGatingTopKHashGenerlized<T>::Process()
|
||||
{
|
||||
CopyInBiasAndInitExpertId();
|
||||
for (int64_t row = 0; row < curCoreRowCount_; row++) {
|
||||
CopyInX(row);
|
||||
ComputeX();
|
||||
if (tilingData_->outFlag) {
|
||||
CopuOutXNorm(row);
|
||||
}
|
||||
SortInGroup();
|
||||
SelectTopKGroupIndex();
|
||||
SelectTopKExpertIdx();
|
||||
SelectTopKExpertScore();
|
||||
CumputeActualTopKExpertId();
|
||||
CopyOut(row);
|
||||
}
|
||||
}
|
||||
} // namespace MoeGatingTopKHash
|
||||
#endif // MOE_GATING_TOP_K_E_K_GENERALIZED_H
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,385 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file moe_gating_top_k_hash_without_group.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef MOE_GATING_TOP_K_E_K_WITHOUT_GROUP_H
|
||||
#define MOE_GATING_TOP_K_E_K_WITHOUT_GROUP_H
|
||||
#include "kernel_operator.h"
|
||||
#include "common.h"
|
||||
#include "kernel_utils.h"
|
||||
namespace MoeGatingTopKHash {
|
||||
using namespace AscendC;
|
||||
|
||||
template <typename T, typename U1, typename U2>
|
||||
class MoeGatingTopKHashWithoutGroup {
|
||||
public:
|
||||
__aicore__ inline MoeGatingTopKHashWithoutGroup(){};
|
||||
__aicore__ inline void Init(GM_ADDR x, GM_ADDR bias, GM_ADDR inputIds, GM_ADDR tid2eid, GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, GM_ADDR workspace,
|
||||
const MoeGatingTopKHashTilingData *tilingData, TPipe *tPipe);
|
||||
__aicore__ inline void Process();
|
||||
|
||||
private:
|
||||
__aicore__ inline void CopyInBiasAndInitExpertId();
|
||||
__aicore__ inline void CopyInX(int64_t progress);
|
||||
__aicore__ inline void ComputeX();
|
||||
__aicore__ inline void CopuOutXNorm(int64_t row);
|
||||
__aicore__ inline void SelectTopKExpertIdx();
|
||||
__aicore__ inline void SelectExpertIdxByHash(int64_t row);
|
||||
__aicore__ inline void SelectTopKExpertScore();
|
||||
__aicore__ inline void CopyOut(int64_t row);
|
||||
|
||||
private:
|
||||
TPipe *pipe_;
|
||||
TQue<QuePosition::VECIN, 1> xInQueue_;
|
||||
TQue<QuePosition::VECOUT, 1> yOutQueue_;
|
||||
TQue<QuePosition::VECOUT, 1> expertIdxOutQueue_;
|
||||
TQue<QuePosition::VECOUT, 1> outOutQueue_;
|
||||
|
||||
TBuf<TPosition::VECCALC> biasBuf_; // 存放输入bias
|
||||
TBuf<TPosition::VECCALC> expertIdBuf_; // 专家编号
|
||||
TBuf<TPosition::VECCALC> xNormWithBiasBuf_; // 存放加了bias之后的值
|
||||
TBuf<TPosition::VECCALC> xNormBuf_; // 存放计算sigmoid或softmax的值
|
||||
TBuf<TPosition::VECCALC> topKExpertIdBuf_;
|
||||
TBuf<TPosition::VECCALC> calcTmpBuf_;
|
||||
|
||||
GlobalTensor<T> xGm_;
|
||||
GlobalTensor<T> biasGm_;
|
||||
GlobalTensor<U1> inputIdsGm_;
|
||||
GlobalTensor<U2> tid2eidGm_;
|
||||
GlobalTensor<T> yGm_;
|
||||
GlobalTensor<int32_t> expertIdxGm_;
|
||||
GlobalTensor<float> outGm_;
|
||||
|
||||
int64_t blockIdx_ = 0;
|
||||
int64_t perCoreRowCount_ = 0;
|
||||
int64_t curCoreRowCount_ = 0;
|
||||
int64_t expertCount_ = 0;
|
||||
bool addBias_ = false;
|
||||
bool outFlag_ = false;
|
||||
bool hashFlag_ = false;
|
||||
int64_t k_ = 0;
|
||||
|
||||
int64_t expertCountAlign_ = 0;
|
||||
const MoeGatingTopKHashTilingData *tilingData_;
|
||||
|
||||
template <HardEvent event>
|
||||
__aicore__ inline void SetWaitFlag(HardEvent evt)
|
||||
{
|
||||
event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(evt));
|
||||
SetFlag<event>(eventId);
|
||||
WaitFlag<event>(eventId);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, typename U1, typename U2>
|
||||
__aicore__ inline void MoeGatingTopKHashWithoutGroup<T, U1, U2>::CopyInBiasAndInitExpertId()
|
||||
{
|
||||
LocalTensor<float> biasTensor = biasBuf_.Get<float>();
|
||||
LocalTensor<int32_t> expertIdTensor = expertIdBuf_.Get<int32_t>();
|
||||
DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(expertCount_ * sizeof(T)), 0, 0, 0};
|
||||
DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast<T>(0)};
|
||||
if (addBias_) {
|
||||
if constexpr (IsSameType<T, float>::value) {
|
||||
DataCopyPad(biasTensor, biasGm_, dataCopyParams, dataCopyPadParams);
|
||||
event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
|
||||
SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
|
||||
WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
|
||||
} else {
|
||||
DataCopyPad(biasTensor[expertCountAlign_].ReinterpretCast<T>(), biasGm_, dataCopyParams, dataCopyPadParams);
|
||||
event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
|
||||
SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
|
||||
WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
|
||||
Cast(biasTensor, biasTensor[expertCountAlign_].ReinterpretCast<T>(), RoundMode::CAST_NONE,
|
||||
expertCountAlign_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
}
|
||||
ArithProgression(expertIdTensor, static_cast<int32_t>(0), static_cast<int32_t>(1), expertCount_);
|
||||
}
|
||||
|
||||
template <typename T, typename U1, typename U2>
|
||||
__aicore__ inline void MoeGatingTopKHashWithoutGroup<T, U1, U2>::CopyInX(int64_t row)
|
||||
{
|
||||
LocalTensor<float> xInLocalTensor = xInQueue_.AllocTensor<float>();
|
||||
DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(expertCount_ * sizeof(T)), 0, 0, 0};
|
||||
DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast<T>(0)};
|
||||
if constexpr (IsSameType<T, float>::value) {
|
||||
DataCopyPad(xInLocalTensor, xGm_[row * expertCount_], dataCopyParams, dataCopyPadParams);
|
||||
} else {
|
||||
DataCopyPad(xInLocalTensor[expertCountAlign_].ReinterpretCast<T>(), xGm_[row * expertCount_], dataCopyParams,
|
||||
dataCopyPadParams);
|
||||
}
|
||||
xInQueue_.EnQue(xInLocalTensor);
|
||||
}
|
||||
|
||||
template <typename T, typename U1, typename U2>
|
||||
__aicore__ inline void MoeGatingTopKHashWithoutGroup<T, U1, U2>::ComputeX()
|
||||
{
|
||||
LocalTensor<float> xNormTensor = xNormBuf_.Get<float>();
|
||||
LocalTensor<float> xInLocalTensor = xInQueue_.DeQue<float>();
|
||||
LocalTensor<float> xNormWithBiasTensor = xNormWithBiasBuf_.Get<float>();
|
||||
LocalTensor<float> biasTensor = biasBuf_.Get<float>();
|
||||
|
||||
if constexpr (!IsSameType<T, float>::value) {
|
||||
Cast(xInLocalTensor, xInLocalTensor[expertCountAlign_].ReinterpretCast<T>(), RoundMode::CAST_NONE,
|
||||
expertCount_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
|
||||
if (tilingData_->normType == NORM_TYPE_SIGMOID) { // sigmoid
|
||||
LocalTensor<uint8_t> calcNormTmpTensor = calcTmpBuf_.Get<uint8_t>();
|
||||
Sigmoid(xNormTensor, xInLocalTensor, calcNormTmpTensor, expertCount_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
} else if (tilingData_->normType == NORM_TYPE_SOFTMAX) { // softmax
|
||||
LocalTensor<float> reduceValueTensor = calcTmpBuf_.Get<float>();
|
||||
LocalTensor<float> calcTmp = calcTmpBuf_.Get<float>()[8];
|
||||
ReduceMax(reduceValueTensor, xInLocalTensor, calcTmp, expertCount_);
|
||||
event_t eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
float maxValue = reduceValueTensor.GetValue(0);
|
||||
event_t eventIdSToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Adds(xNormTensor, xInLocalTensor, -maxValue, expertCount_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Exp(xNormTensor, xNormTensor, expertCount_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
ReduceSum(reduceValueTensor, xNormTensor, calcTmp, expertCount_);
|
||||
eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
float sumValue = reduceValueTensor.GetValue(0);
|
||||
eventIdSToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Muls(xNormTensor, xNormTensor, 1.0f / sumValue, expertCount_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
} else {
|
||||
LocalTensor<float> calcNormTmpTensor = calcTmpBuf_.Get<float>();
|
||||
Exp(calcNormTmpTensor, xInLocalTensor, expertCount_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Adds(calcNormTmpTensor, calcNormTmpTensor, float(1.0), expertCount_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Ln(calcNormTmpTensor, calcNormTmpTensor, expertCount_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Sqrt(xNormTensor, calcNormTmpTensor, expertCount_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
if (addBias_) {
|
||||
Add(xNormWithBiasTensor, xNormTensor, biasTensor, expertCount_);
|
||||
} else {
|
||||
DataCopy(xNormWithBiasTensor, xNormTensor, expertCountAlign_);
|
||||
}
|
||||
|
||||
int64_t duplicateNum = expertCount_ % ONE_REPEAT_SORT_NUM;
|
||||
int duplicateIndex = expertCount_ - duplicateNum;
|
||||
if (duplicateNum > 0) {
|
||||
uint64_t mask0 = UINT64_MAX;
|
||||
mask0 = mask0 << duplicateNum;
|
||||
mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
|
||||
uint64_t mask[2] = {mask0, 0};
|
||||
Duplicate(xNormWithBiasTensor.ReinterpretCast<int32_t>()[duplicateIndex], FLOAT32_NEG_INF, mask, 1, 1, 1);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
xInQueue_.FreeTensor(xInLocalTensor);
|
||||
}
|
||||
|
||||
template <typename T, typename U1, typename U2>
|
||||
__aicore__ inline void MoeGatingTopKHashWithoutGroup<T, U1, U2>::CopuOutXNorm(int64_t row)
|
||||
{
|
||||
LocalTensor<float> outOutTensor = outOutQueue_.AllocTensor<float>();
|
||||
LocalTensor<float> xNormTensor = xNormBuf_.Get<float>();
|
||||
DataCopy(outOutTensor, xNormTensor, expertCountAlign_);
|
||||
outOutQueue_.EnQue<float>(outOutTensor);
|
||||
outOutTensor = outOutQueue_.DeQue<float>();
|
||||
DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(expertCount_ * sizeof(float)), 0, 0, 0};
|
||||
DataCopyPad(outGm_[row * expertCount_], outOutTensor, dataCopyParams);
|
||||
outOutQueue_.FreeTensor(outOutTensor);
|
||||
}
|
||||
|
||||
template <typename T, typename U1, typename U2>
|
||||
__aicore__ inline void MoeGatingTopKHashWithoutGroup<T, U1, U2>::SelectTopKExpertIdx()
|
||||
{
|
||||
LocalTensor<int32_t> expertIdxOut = expertIdxOutQueue_.AllocTensor<int32_t>();
|
||||
LocalTensor<float> xNormWithBiasTensor = xNormWithBiasBuf_.Get<float>();
|
||||
LocalTensor<uint32_t> expertIdTensor = expertIdBuf_.Get<uint32_t>();
|
||||
LocalTensor<int32_t> topKExpertId = topKExpertIdBuf_.Get<int32_t>();
|
||||
LocalTensor<float> sortedScore = calcTmpBuf_.Get<float>();
|
||||
LocalTensor<float> sortTmp = calcTmpBuf_.Get<float>()[expertCountAlign_ * CONSTANT_TWO];
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
Sort<float, true>(sortedScore, xNormWithBiasTensor, expertIdTensor, sortTmp,
|
||||
expertCountAlign_ / ONE_REPEAT_SORT_NUM);
|
||||
|
||||
GatherMaskParams gatherMaskParams;
|
||||
gatherMaskParams.repeatTimes = Ceil(k_ * sizeof(float) * CONSTANT_TWO, REPEAT_BYTES);
|
||||
gatherMaskParams.src0BlockStride = 1;
|
||||
gatherMaskParams.src0RepeatStride = REPEAT_BLOCKS;
|
||||
gatherMaskParams.src1RepeatStride = 0;
|
||||
|
||||
uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数
|
||||
uint8_t src1Pattern = 2; // 内置固定模式
|
||||
PipeBarrier<PIPE_V>();
|
||||
GatherMask(topKExpertId, sortedScore.template ReinterpretCast<int32_t>(), src1Pattern, false,
|
||||
static_cast<uint32_t>(0), gatherMaskParams, rsvdCnt);
|
||||
|
||||
DataCopy(expertIdxOut, topKExpertId, expertCountAlign_);
|
||||
expertIdxOutQueue_.EnQue<int32_t>(expertIdxOut);
|
||||
}
|
||||
|
||||
template <typename T, typename U1, typename U2>
|
||||
__aicore__ inline void MoeGatingTopKHashWithoutGroup<T, U1, U2>::SelectTopKExpertScore()
|
||||
{
|
||||
LocalTensor<float> xNormTensor = xNormBuf_.Get<float>();
|
||||
LocalTensor<float> yOutTensor = yOutQueue_.AllocTensor<float>();
|
||||
LocalTensor<int32_t> topKExpertId = topKExpertIdBuf_.Get<int32_t>();
|
||||
LocalTensor<int32_t> topKExpertIdWithByte = calcTmpBuf_.Get<int32_t>();
|
||||
PipeBarrier<PIPE_V>();
|
||||
Muls(topKExpertIdWithByte, topKExpertId, static_cast<int32_t>(sizeof(float)), k_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Gather(yOutTensor, xNormTensor, topKExpertIdWithByte.template ReinterpretCast<uint32_t>(), static_cast<uint32_t>(0),
|
||||
k_);
|
||||
|
||||
if (tilingData_->normType == NORM_TYPE_SIGMOID || tilingData_->normType == NORM_TYPE_SOFTPLUS) {
|
||||
LocalTensor<float> maxValueTensor = calcTmpBuf_.Get<float>();
|
||||
LocalTensor<float> tmpTensor = calcTmpBuf_.Get<float>()[BLOCK_BYTES];
|
||||
PipeBarrier<PIPE_V>();
|
||||
ReduceSum(maxValueTensor, yOutTensor, tmpTensor, k_);
|
||||
event_t eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
float sumValue = maxValueTensor.GetValue(0) + tilingData_->eps;
|
||||
event_t eventIdSToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Duplicate(tmpTensor, sumValue, k_);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Div(yOutTensor, yOutTensor, tmpTensor, k_);
|
||||
}
|
||||
PipeBarrier<PIPE_V>();
|
||||
Muls(yOutTensor, yOutTensor, tilingData_->routedScalingFactor, k_);
|
||||
|
||||
if constexpr (!IsSameType<T, float>::value) {
|
||||
PipeBarrier<PIPE_V>();
|
||||
Cast(yOutTensor.ReinterpretCast<T>(), yOutTensor, RoundMode::CAST_RINT, k_);
|
||||
}
|
||||
|
||||
yOutQueue_.EnQue<float>(yOutTensor);
|
||||
}
|
||||
|
||||
template <typename T, typename U1, typename U2>
|
||||
__aicore__ inline void MoeGatingTopKHashWithoutGroup<T, U1, U2>::CopyOut(int64_t row)
|
||||
{
|
||||
LocalTensor<T> yOutTensor = yOutQueue_.DeQue<T>();
|
||||
LocalTensor<int32_t> expertIdxOut = expertIdxOutQueue_.DeQue<int32_t>();
|
||||
DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(k_ * sizeof(T)), 0, 0, 0};
|
||||
DataCopyPad(yGm_[row * k_], yOutTensor, dataCopyParams);
|
||||
dataCopyParams.blockLen = k_ * sizeof(int32_t);
|
||||
DataCopyPad(expertIdxGm_[row * k_], expertIdxOut, dataCopyParams);
|
||||
yOutQueue_.FreeTensor(yOutTensor);
|
||||
expertIdxOutQueue_.FreeTensor(expertIdxOut);
|
||||
}
|
||||
|
||||
template <typename T, typename U1, typename U2>
|
||||
__aicore__ inline void MoeGatingTopKHashWithoutGroup<T, U1, U2>::SelectExpertIdxByHash(int64_t row)
|
||||
{
|
||||
LocalTensor<int32_t> expertIdxOut = expertIdxOutQueue_.AllocTensor<int32_t>();
|
||||
LocalTensor<U2> hashExpertId = topKExpertIdBuf_.Get<U2>();
|
||||
LocalTensor<int32_t> hashExpertIdInt32 = hashExpertId.template ReinterpretCast<int32_t>();
|
||||
U1 key = inputIdsGm_.GetValue(row);
|
||||
SetWaitFlag<HardEvent::S_MTE2>(HardEvent::S_MTE2);
|
||||
DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(k_ * sizeof(U2)), 0, 0, 0};
|
||||
DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast<U2>(0)};
|
||||
DataCopyPad(hashExpertId, tid2eidGm_[key * k_], dataCopyParams, dataCopyPadParams);
|
||||
SetWaitFlag<HardEvent::MTE2_V>(HardEvent::MTE2_V);
|
||||
if constexpr (IsSameType<U2, int32_t>::value) {
|
||||
DataCopy(expertIdxOut, hashExpertId, Align(k_, sizeof(int32_t)));
|
||||
} else {
|
||||
Cast(hashExpertIdInt32, hashExpertId, RoundMode::CAST_NONE, Align(k_, sizeof(U2)));
|
||||
PipeBarrier<PIPE_V>();
|
||||
DataCopy(expertIdxOut, hashExpertIdInt32, Align(k_, sizeof(int32_t)));
|
||||
}
|
||||
expertIdxOutQueue_.EnQue<int32_t>(expertIdxOut);
|
||||
}
|
||||
|
||||
template <typename T, typename U1, typename U2>
|
||||
__aicore__ inline void MoeGatingTopKHashWithoutGroup<T, U1, U2>::Init(GM_ADDR x, GM_ADDR bias, GM_ADDR inputIds, GM_ADDR tid2eid,
|
||||
GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, GM_ADDR workspace,
|
||||
const MoeGatingTopKHashTilingData *tilingData, TPipe *tPipe)
|
||||
{
|
||||
tilingData_ = tilingData;
|
||||
pipe_ = tPipe;
|
||||
blockIdx_ = GetBlockIdx();
|
||||
perCoreRowCount_ = tilingData_->perCoreRowCount;
|
||||
if (blockIdx_ == GetBlockNum() - 1) {
|
||||
curCoreRowCount_ = tilingData_->lastCoreRowCount;
|
||||
} else {
|
||||
curCoreRowCount_ = tilingData_->perCoreRowCount;
|
||||
}
|
||||
expertCount_ = tilingData_->expertCount;
|
||||
addBias_ = tilingData_->addBias == 1;
|
||||
outFlag_ = tilingData_->outFlag == 1;
|
||||
hashFlag_ = tilingData_->hashFlag == 1;
|
||||
k_ = tilingData_->k;
|
||||
|
||||
expertCountAlign_ = Ceil(expertCount_, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
|
||||
|
||||
// init input gm buf
|
||||
xGm_.SetGlobalBuffer((__gm__ T *)x + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_);
|
||||
biasGm_.SetGlobalBuffer((__gm__ T *)bias, expertCount_);
|
||||
inputIdsGm_.SetGlobalBuffer((__gm__ U1 *)inputIds);
|
||||
tid2eidGm_.SetGlobalBuffer((__gm__ U2 *)tid2eid);
|
||||
|
||||
// init output gm buf
|
||||
yGm_.SetGlobalBuffer((__gm__ T *)y + perCoreRowCount_ * k_ * blockIdx_, k_);
|
||||
expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx + perCoreRowCount_ * k_ * blockIdx_, k_);
|
||||
outGm_.SetGlobalBuffer((__gm__ float *)out + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_);
|
||||
|
||||
// init que
|
||||
pipe_->InitBuffer(xInQueue_, 1, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T)));
|
||||
pipe_->InitBuffer(yOutQueue_, 1, Align(k_, sizeof(float)) * sizeof(float));
|
||||
pipe_->InitBuffer(expertIdxOutQueue_, 1, Align(k_, sizeof(float)) * sizeof(int32_t));
|
||||
pipe_->InitBuffer(outOutQueue_, 1, expertCountAlign_ * sizeof(float));
|
||||
|
||||
// init calc buf
|
||||
pipe_->InitBuffer(biasBuf_, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T)));
|
||||
pipe_->InitBuffer(expertIdBuf_, expertCountAlign_ * sizeof(int32_t));
|
||||
pipe_->InitBuffer(xNormBuf_, expertCountAlign_ * sizeof(float));
|
||||
pipe_->InitBuffer(xNormWithBiasBuf_, expertCountAlign_ * sizeof(float));
|
||||
pipe_->InitBuffer(topKExpertIdBuf_, Align(k_, sizeof(U2)) * sizeof(U2));
|
||||
|
||||
// init tmp buf
|
||||
pipe_->InitBuffer(calcTmpBuf_, expertCountAlign_ * sizeof(float) * CONSTANT_EIGHT);
|
||||
}
|
||||
|
||||
template <typename T, typename U1, typename U2>
|
||||
__aicore__ inline void MoeGatingTopKHashWithoutGroup<T, U1, U2>::Process()
|
||||
{
|
||||
CopyInBiasAndInitExpertId();
|
||||
for (int64_t row = 0; row < curCoreRowCount_; row++) {
|
||||
CopyInX(row);
|
||||
ComputeX();
|
||||
if (outFlag_) {
|
||||
CopuOutXNorm(row);
|
||||
}
|
||||
if (hashFlag_) {
|
||||
SelectExpertIdxByHash(row + perCoreRowCount_ * blockIdx_);
|
||||
} else {
|
||||
SelectTopKExpertIdx();
|
||||
}
|
||||
SelectTopKExpertScore();
|
||||
CopyOut(row);
|
||||
}
|
||||
}
|
||||
} // namespace MoeGatingTopKHash
|
||||
#endif // MOE_GATING_TOP_K_E_K_WITHOUT_GROUP_H
|
||||
Reference in New Issue
Block a user