22
csrc/moe/causal_conv1d_v310/op_host/CMakeLists.txt
Normal file
22
csrc/moe/causal_conv1d_v310/op_host/CMakeLists.txt
Normal file
@@ -0,0 +1,22 @@
|
||||
add_op_to_compiled_list()
|
||||
|
||||
if (BUILD_OPEN_PROJECT)
|
||||
target_sources(op_host_aclnn PRIVATE
|
||||
causal_conv1d_v310_def.cpp
|
||||
)
|
||||
endif()
|
||||
|
||||
add_ops_compile_options(
|
||||
OP_NAME CausalConv1dFnV310
|
||||
OPTIONS
|
||||
--cce-auto-sync=on
|
||||
-Wno-deprecated-declarations
|
||||
)
|
||||
|
||||
if (NOT BUILD_OPS_RTY_KERNEL)
|
||||
add_modules_sources(OPTYPE causal_conv1d_v310 ACLNNTYPE aclnn)
|
||||
target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE
|
||||
${CMAKE_CURRENT_SOURCE_DIR}
|
||||
)
|
||||
endif()
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
/**
|
||||
* This program is free software, you can redistribute it and/or modify it.
|
||||
* Copyright (c) 2026 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.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file causal_conv1d_v310_def.cpp
|
||||
* \brief
|
||||
*/
|
||||
#include "register/op_def_registry.h"
|
||||
|
||||
namespace ops {
|
||||
|
||||
class CausalConv1dV310 : public OpDef {
|
||||
public:
|
||||
explicit CausalConv1dV310(const char *name) : OpDef(name)
|
||||
{
|
||||
this->Input("x")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT16})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("weight")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT16})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("bias")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataType({ge::DT_FLOAT16})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("convStates")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT16})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("queryStartLoc")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataTypeList({ge::DT_INT32, ge::DT_INT64})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("cacheIndices")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataTypeList({ge::DT_INT32, ge::DT_INT64})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("initialStateMode")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataTypeList({ge::DT_BOOL, ge::DT_INT32, ge::DT_INT64})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("numAcceptedTokens")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataTypeList({ge::DT_INT32, ge::DT_INT64})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
|
||||
this->Output("y")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT16})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
|
||||
this->Attr("activationMode").AttrType(OPTIONAL).Int(0);
|
||||
this->Attr("padSlotId").AttrType(OPTIONAL).Int(-1);
|
||||
this->Attr("runMode").AttrType(OPTIONAL).Int(0);
|
||||
|
||||
OpAICoreConfig aicoreConfig;
|
||||
aicoreConfig.DynamicCompileStaticFlag(true)
|
||||
.DynamicFormatFlag(false)
|
||||
.DynamicRankSupportFlag(true)
|
||||
.DynamicShapeSupportFlag(true)
|
||||
.NeedCheckSupportFlag(false)
|
||||
.PrecisionReduceFlag(true)
|
||||
.ExtendCfgInfo("coreType.value", "AiCore");
|
||||
this->AICore().AddConfig("ascend310p", aicoreConfig);
|
||||
}
|
||||
};
|
||||
OP_ADD(CausalConv1dV310);
|
||||
|
||||
} // namespace ops
|
||||
@@ -0,0 +1,41 @@
|
||||
/**
|
||||
* This program is free software, you can redistribute it and/or modify it.
|
||||
* Copyright (c) 2026 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.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file causal_conv1d_v310_infershape.cpp
|
||||
* \brief
|
||||
*/
|
||||
#include "register/op_impl_registry.h"
|
||||
#include "tiling_base/error_log.h"
|
||||
|
||||
using namespace ge;
|
||||
|
||||
namespace ops {
|
||||
static constexpr int64_t IDX_0 = 0;
|
||||
|
||||
static ge::graphStatus InferShapeCausalConv1d(gert::InferShapeContext *context)
|
||||
{
|
||||
OP_LOGD(context->GetNodeName(), "Begin to do InferShapeCausalConv1d");
|
||||
|
||||
const gert::Shape *xShape = context->GetInputShape(IDX_0);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, xShape);
|
||||
|
||||
gert::Shape *yShape = context->GetOutputShape(IDX_0);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, yShape);
|
||||
|
||||
*yShape = *xShape;
|
||||
|
||||
OP_LOGD(context->GetNodeName(), "End to do InferShapeCausalConv1d");
|
||||
return GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
IMPL_OP_INFERSHAPE(CausalConv1dV310).InferShape(InferShapeCausalConv1d);
|
||||
} // namespace ops
|
||||
@@ -0,0 +1,521 @@
|
||||
/**
|
||||
* This program is free software, you can redistribute it and/or modify it.
|
||||
* Copyright (c) 2026 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.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file causal_conv1d_v310_tiling.cpp
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#include "math_util.h"
|
||||
#include "tiling_base/tiling_templates_registry.h"
|
||||
#include "tiling_base/tiling_util.h"
|
||||
#include "../op_kernel/causal_conv1d_v310_tiling_data.h"
|
||||
#include "../op_kernel/causal_conv1d_v310_tiling_key.h"
|
||||
|
||||
#include <set>
|
||||
#include <limits>
|
||||
|
||||
namespace optiling {
|
||||
|
||||
using namespace Ops::Transformer::OpTiling;
|
||||
|
||||
constexpr uint32_t X_INDEX = 0;
|
||||
constexpr uint32_t WEIGHT_INDEX = 1;
|
||||
constexpr uint32_t BIAS_INDEX = 2;
|
||||
constexpr uint32_t CONV_STATES_INDEX = 3;
|
||||
constexpr uint32_t QUERY_START_LOC_INDEX = 4;
|
||||
constexpr uint32_t CACHE_INDICES_INDEX = 5;
|
||||
constexpr uint32_t INITIAL_STATE_MODE_INDEX = 6;
|
||||
constexpr uint32_t NUM_ACCEPTED_TOKENS_INDEX = 7;
|
||||
|
||||
constexpr int32_t ATTR_ACTIVATION_MODE_INDEX = 0;
|
||||
constexpr int32_t ATTR_PAD_SLOT_ID_INDEX = 1;
|
||||
constexpr int32_t ATTR_RUN_MODE_INDEX = 2;
|
||||
|
||||
struct CausalConv1dCompileInfo {
|
||||
uint64_t ubSize = 0;
|
||||
uint32_t coreNum = 0;
|
||||
};
|
||||
|
||||
struct DimTileChoice {
|
||||
int64_t dimTileSize = 0;
|
||||
int64_t blocksPerSeq = 0;
|
||||
int64_t gridSize = 0;
|
||||
};
|
||||
|
||||
static inline int64_t CeilDivInt64(int64_t x, int64_t y)
|
||||
{
|
||||
return (x + y - 1) / y;
|
||||
}
|
||||
|
||||
static inline bool FitsInInt32(int64_t v)
|
||||
{
|
||||
return v >= static_cast<int64_t>(std::numeric_limits<int32_t>::min()) &&
|
||||
v <= static_cast<int64_t>(std::numeric_limits<int32_t>::max());
|
||||
}
|
||||
|
||||
static inline DimTileChoice ChooseDimTileSize(gert::TilingContext *context, int64_t batch, int64_t dim,
|
||||
uint32_t coreNum)
|
||||
{
|
||||
const int64_t candidates[] = {4096, 2048, 1024, 512, 384, 192};
|
||||
|
||||
auto ChooseOnce = [&](bool requireExactDiv) -> DimTileChoice {
|
||||
const bool preferMoreBlocks = (batch > static_cast<int64_t>(coreNum));
|
||||
DimTileChoice bestOver;
|
||||
int64_t bestOverGap = std::numeric_limits<int64_t>::max();
|
||||
DimTileChoice bestUnder;
|
||||
|
||||
for (int64_t dimTileSize : candidates) {
|
||||
if (dimTileSize <= 0) {
|
||||
continue;
|
||||
}
|
||||
if (requireExactDiv && (dim % dimTileSize != 0)) {
|
||||
continue;
|
||||
}
|
||||
const int64_t testBlocksPerSeq = requireExactDiv ? (dim / dimTileSize) : CeilDivInt64(dim, dimTileSize);
|
||||
if (preferMoreBlocks && testBlocksPerSeq <= 1) {
|
||||
continue;
|
||||
}
|
||||
const int64_t blocksPerSeq = requireExactDiv ? (dim / dimTileSize) : CeilDivInt64(dim, dimTileSize);
|
||||
const int64_t gridSize = batch * blocksPerSeq;
|
||||
if (gridSize <= 0) {
|
||||
continue;
|
||||
}
|
||||
OP_LOGD(context, "DimTile candidate[%s]: dimTileSize[%ld], blocksPerSeq[%ld], gridSize[%ld], coreNum[%u].",
|
||||
requireExactDiv ? "exact" : "tail", dimTileSize, blocksPerSeq, gridSize, coreNum);
|
||||
if (gridSize >= static_cast<int64_t>(coreNum)) {
|
||||
const int64_t gap = gridSize - static_cast<int64_t>(coreNum);
|
||||
if (gap < bestOverGap) {
|
||||
bestOver.dimTileSize = dimTileSize;
|
||||
bestOver.blocksPerSeq = blocksPerSeq;
|
||||
bestOver.gridSize = gridSize;
|
||||
bestOverGap = gap;
|
||||
}
|
||||
} else if (gridSize > bestUnder.gridSize ||
|
||||
(gridSize == bestUnder.gridSize && dimTileSize < bestUnder.dimTileSize)) {
|
||||
bestUnder.dimTileSize = dimTileSize;
|
||||
bestUnder.blocksPerSeq = blocksPerSeq;
|
||||
bestUnder.gridSize = gridSize;
|
||||
}
|
||||
}
|
||||
return (bestOver.dimTileSize != 0) ? bestOver : bestUnder;
|
||||
};
|
||||
|
||||
DimTileChoice result = ChooseOnce(true /*requireExactDiv*/);
|
||||
if (result.dimTileSize == 0) {
|
||||
result = ChooseOnce(false /*requireExactDiv*/);
|
||||
}
|
||||
OP_LOGD(context, "DimTile chosen: dimTileSize[%ld], blocksPerSeq[%ld], gridSize[%ld].", result.dimTileSize,
|
||||
result.blocksPerSeq, result.gridSize);
|
||||
return result;
|
||||
}
|
||||
|
||||
static ge::graphStatus GetPlatformInfo(gert::TilingContext *context, uint64_t &ubSize, uint32_t &coreNum)
|
||||
{
|
||||
fe::PlatFormInfos *platformInfoPtr = context->GetPlatformInfo();
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, platformInfoPtr);
|
||||
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr);
|
||||
coreNum = ascendcPlatform.GetCoreNumAiv();
|
||||
OP_CHECK_IF(coreNum == 0, OP_LOGE(context, "coreNum is 0"), return ge::GRAPH_FAILED);
|
||||
ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
|
||||
OP_CHECK_IF(ubSize == 0, OP_LOGE(context, "ubSize is 0"), return ge::GRAPH_FAILED);
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
static ge::graphStatus GetWorkspaceSize(gert::TilingContext *context)
|
||||
{
|
||||
size_t *currentWorkspace = context->GetWorkspaceSizes(1);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, currentWorkspace);
|
||||
currentWorkspace[0] = 0;
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
static ge::graphStatus GetAttrsInfo(gert::TilingContext *context, int64_t &activationMode, int64_t &padSlotId,
|
||||
int64_t &runMode)
|
||||
{
|
||||
auto attrs = context->GetAttrs();
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
|
||||
|
||||
const int64_t *activationModePtr = attrs->GetAttrPointer<int64_t>(ATTR_ACTIVATION_MODE_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, activationModePtr);
|
||||
activationMode = *activationModePtr;
|
||||
OP_CHECK_IF(activationMode != 0 && activationMode != 1, OP_LOGE(context, "activationMode only supports 0/1"),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
const int64_t *padSlotIdPtr = attrs->GetAttrPointer<int64_t>(ATTR_PAD_SLOT_ID_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, padSlotIdPtr);
|
||||
padSlotId = *padSlotIdPtr;
|
||||
|
||||
const int64_t *runModePtr = attrs->GetAttrPointer<int64_t>(ATTR_RUN_MODE_INDEX);
|
||||
runMode = (runModePtr == nullptr) ? 0 : *runModePtr;
|
||||
OP_CHECK_IF(runMode != 0 && runMode != 1, OP_LOGE(context, "runMode only supports 0/1"), return ge::GRAPH_FAILED);
|
||||
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
static ge::graphStatus GetShapeDtypeInfo(gert::TilingContext *context, CausalConv1dTilingData &tiling)
|
||||
{
|
||||
const bool isDecodeMode = (tiling.runMode == 1);
|
||||
|
||||
auto xShapePtr = context->GetInputShape(X_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, xShapePtr);
|
||||
auto xShape = EnsureNotScalar(xShapePtr->GetStorageShape());
|
||||
|
||||
int64_t dim = 0;
|
||||
int64_t cuSeqlen = 0;
|
||||
int64_t seqLen = 0;
|
||||
int64_t batch = 0;
|
||||
int64_t inputMode = 0;
|
||||
|
||||
if (xShape.GetDimNum() == 2) {
|
||||
if (isDecodeMode) {
|
||||
inputMode = 2;
|
||||
batch = xShape.GetDim(0);
|
||||
dim = xShape.GetDim(1);
|
||||
seqLen = 1;
|
||||
cuSeqlen = batch;
|
||||
OP_CHECK_IF(batch <= 0 || dim <= 0, OP_LOGE(context, "invalid x shape for 2D decode mode"),
|
||||
return ge::GRAPH_FAILED);
|
||||
} else {
|
||||
inputMode = 0;
|
||||
cuSeqlen = xShape.GetDim(0);
|
||||
dim = xShape.GetDim(1);
|
||||
seqLen = 0;
|
||||
OP_CHECK_IF(dim <= 0 || cuSeqlen < 0, OP_LOGE(context, "invalid x shape for 2D varlen mode"),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
} else if (xShape.GetDimNum() == 3) {
|
||||
inputMode = 1;
|
||||
batch = xShape.GetDim(0);
|
||||
seqLen = xShape.GetDim(1);
|
||||
dim = xShape.GetDim(2);
|
||||
cuSeqlen = batch * seqLen;
|
||||
OP_CHECK_IF(batch <= 0 || dim <= 0 || seqLen <= 0, OP_LOGE(context, "invalid x shape for 3D batch mode"),
|
||||
return ge::GRAPH_FAILED);
|
||||
} else {
|
||||
OP_LOGE(context, "x must be 2D (cu_seqlen, dim) or 3D (batch, seqlen, dim)");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
|
||||
auto wShapePtr = context->GetInputShape(WEIGHT_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, wShapePtr);
|
||||
auto wShape = EnsureNotScalar(wShapePtr->GetStorageShape());
|
||||
OP_CHECK_IF(wShape.GetDimNum() != 2, OP_LOGE(context, "weight must be 2D: (width, dim)"), return ge::GRAPH_FAILED);
|
||||
const int64_t width = wShape.GetDim(0);
|
||||
const int64_t wDim = wShape.GetDim(1);
|
||||
OP_CHECK_IF(wDim != dim, OP_LOGE(context, "weight.shape[1] must equal dim"), return ge::GRAPH_FAILED);
|
||||
OP_CHECK_IF(width < 2 || width > 4, OP_LOGE(context, "Only support width in [2,4] now, actually is %ld.", width),
|
||||
return ge::GRAPH_FAILED);
|
||||
OP_CHECK_IF(dim % 16 != 0,
|
||||
OP_LOGE(context, "dim must be a multiple of 16 for fp16/bf16 alignment, actually is %ld.", dim),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
auto sShapePtr = context->GetInputShape(CONV_STATES_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, sShapePtr);
|
||||
auto sShape = EnsureNotScalar(sShapePtr->GetStorageShape());
|
||||
OP_CHECK_IF(sShape.GetDimNum() != 3, OP_LOGE(context, "convStates must be 3D: (num_cache_lines, state_len, dim)"),
|
||||
return ge::GRAPH_FAILED);
|
||||
const int64_t numCacheLines = sShape.GetDim(0);
|
||||
const int64_t stateLen = sShape.GetDim(1);
|
||||
const int64_t sDim = sShape.GetDim(2);
|
||||
OP_CHECK_IF(numCacheLines <= 0, OP_LOGE(context, "convStates.shape[0] (num_cache_lines) must be > 0"),
|
||||
return ge::GRAPH_FAILED);
|
||||
OP_CHECK_IF(sDim != dim, OP_LOGE(context, "convStates.shape[2] must equal dim"), return ge::GRAPH_FAILED);
|
||||
OP_CHECK_IF(stateLen < (width - 1), OP_LOGE(context, "convStates.shape[1] must be >= width-1"),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
auto qslShapePtr = context->GetOptionalInputShape(QUERY_START_LOC_INDEX);
|
||||
const gert::CompileTimeTensorDesc *qslDesc = context->GetOptionalInputDesc(QUERY_START_LOC_INDEX);
|
||||
bool qslAbsent = true;
|
||||
int64_t qslSize = 0;
|
||||
if (qslShapePtr != nullptr) {
|
||||
const auto qslStorageShape = qslShapePtr->GetStorageShape();
|
||||
const int64_t qslDimNum = qslStorageShape.GetDimNum();
|
||||
qslAbsent = (qslDimNum == 0) || (qslDimNum == 1 && qslStorageShape.GetDim(0) <= 0);
|
||||
|
||||
if (!qslAbsent) {
|
||||
auto qslShape = EnsureNotScalar(qslStorageShape);
|
||||
OP_CHECK_IF(qslShape.GetDimNum() != 1, OP_LOGE(context, "queryStartLoc must be 1D"),
|
||||
return ge::GRAPH_FAILED);
|
||||
qslSize = qslShape.GetDim(0);
|
||||
OP_CHECK_IF(qslSize < 1, OP_LOGE(context, "queryStartLoc.size must be >= 1"), return ge::GRAPH_FAILED);
|
||||
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, qslDesc);
|
||||
const ge::DataType qslDtype = qslDesc->GetDataType();
|
||||
OP_CHECK_IF(qslDtype != ge::DT_INT32 && qslDtype != ge::DT_INT64,
|
||||
OP_LOGE(context, "queryStartLoc dtype must be int32 or int64"), return ge::GRAPH_FAILED);
|
||||
}
|
||||
}
|
||||
|
||||
if (qslAbsent) {
|
||||
OP_CHECK_IF(inputMode == 0, OP_LOGE(context, "queryStartLoc is required in 2D varlen mode (inputMode=0)"),
|
||||
return ge::GRAPH_FAILED);
|
||||
qslSize = batch + 1;
|
||||
}
|
||||
tiling.hasQueryStartLoc = qslAbsent ? 0 : 1;
|
||||
|
||||
OP_CHECK_IF(cuSeqlen > static_cast<int64_t>(std::numeric_limits<int32_t>::max()),
|
||||
OP_LOGE(context, "cuSeqlen is too large for int32 indexing, got %ld", cuSeqlen),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
if (!qslAbsent && isDecodeMode && inputMode == 2) {
|
||||
const int64_t batchFromQsl = qslSize - 1;
|
||||
if (batchFromQsl != batch) {
|
||||
inputMode = 0;
|
||||
cuSeqlen = xShape.GetDim(0);
|
||||
batch = batchFromQsl;
|
||||
seqLen = 0;
|
||||
OP_CHECK_IF(dim <= 0 || cuSeqlen < 0 || batch < 0,
|
||||
OP_LOGE(context, "invalid x/queryStartLoc shapes for 2D varlen decode mode"),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
}
|
||||
|
||||
if (inputMode == 0) {
|
||||
batch = qslSize - 1;
|
||||
}
|
||||
|
||||
if (!qslAbsent && (inputMode == 1 || inputMode == 2)) {
|
||||
OP_CHECK_IF(qslSize != batch + 1, OP_LOGE(context, "queryStartLoc.size must equal batch + 1"),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
|
||||
if (isDecodeMode) {
|
||||
const int64_t decodeSeqLen = (inputMode == 1) ? seqLen : 1;
|
||||
OP_CHECK_IF(decodeSeqLen < 1, OP_LOGE(context, "decode mode requires seqlen >= 1, actual is %ld", decodeSeqLen),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
|
||||
tiling.hasCacheIndices = 0;
|
||||
tiling.cacheIndicesStride = 1;
|
||||
bool ciAbsent = true;
|
||||
auto ciShapePtr = context->GetOptionalInputShape(CACHE_INDICES_INDEX);
|
||||
if (ciShapePtr != nullptr) {
|
||||
const auto ciStorageShape = ciShapePtr->GetStorageShape();
|
||||
const int64_t ciDimNum = ciStorageShape.GetDimNum();
|
||||
ciAbsent = (ciDimNum == 0) || (ciDimNum == 1 && ciStorageShape.GetDim(0) <= 0);
|
||||
if (!ciAbsent) {
|
||||
auto ciShape = EnsureNotScalar(ciStorageShape);
|
||||
// Spec decode passes cache indices as [batch, num_spec + 1];
|
||||
// kernels read the first column with cacheIndicesStride.
|
||||
OP_CHECK_IF(ciShape.GetDimNum() != 1 && ciShape.GetDimNum() != 2,
|
||||
OP_LOGE(context, "cacheIndices must be 1D or 2D"), return ge::GRAPH_FAILED);
|
||||
OP_CHECK_IF(ciShape.GetDim(0) != batch, OP_LOGE(context, "cacheIndices first dim must equal batch"),
|
||||
return ge::GRAPH_FAILED);
|
||||
if (ciShape.GetDimNum() == 2) {
|
||||
OP_CHECK_IF(ciShape.GetDim(1) <= 0, OP_LOGE(context, "cacheIndices second dim must be positive"),
|
||||
return ge::GRAPH_FAILED);
|
||||
tiling.cacheIndicesStride = ciShape.GetDim(1);
|
||||
}
|
||||
tiling.hasCacheIndices = 1;
|
||||
}
|
||||
}
|
||||
if (ciAbsent) {
|
||||
OP_CHECK_IF(numCacheLines < batch,
|
||||
OP_LOGE(context,
|
||||
"cacheIndices is absent, requires convStates.shape[0] (num_cache_lines) >= batch for "
|
||||
"identity mapping, got num_cache_lines=%ld batch=%ld",
|
||||
numCacheLines, batch),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
|
||||
tiling.hasInitialStateMode = 0;
|
||||
auto ismShapePtr = context->GetOptionalInputShape(INITIAL_STATE_MODE_INDEX);
|
||||
if (ismShapePtr != nullptr) {
|
||||
const auto ismStorageShape = ismShapePtr->GetStorageShape();
|
||||
const int64_t ismDimNum = ismStorageShape.GetDimNum();
|
||||
const bool ismAbsent = (ismDimNum == 0) || (ismDimNum == 1 && ismStorageShape.GetDim(0) <= 0);
|
||||
if (!ismAbsent) {
|
||||
auto ismShape = EnsureNotScalar(ismStorageShape);
|
||||
OP_CHECK_IF(ismShape.GetDimNum() != 1, OP_LOGE(context, "initialStateMode must be 1D"),
|
||||
return ge::GRAPH_FAILED);
|
||||
OP_CHECK_IF(ismShape.GetDim(0) != batch, OP_LOGE(context, "initialStateMode.size must equal batch"),
|
||||
return ge::GRAPH_FAILED);
|
||||
tiling.hasInitialStateMode = 1;
|
||||
}
|
||||
}
|
||||
|
||||
tiling.hasNumAcceptedTokens = 0;
|
||||
auto natShapePtr = context->GetOptionalInputShape(NUM_ACCEPTED_TOKENS_INDEX);
|
||||
if (natShapePtr != nullptr) {
|
||||
const auto natStorageShape = natShapePtr->GetStorageShape();
|
||||
const int64_t natDimNum = natStorageShape.GetDimNum();
|
||||
const bool natAbsent = (natDimNum == 0) || (natDimNum == 1 && natStorageShape.GetDim(0) <= 0);
|
||||
if (!natAbsent) {
|
||||
OP_CHECK_IF(!isDecodeMode,
|
||||
OP_LOGE(context, "numAcceptedTokens is only supported in runMode=1 (decode/update)"),
|
||||
return ge::GRAPH_FAILED);
|
||||
auto natShape = EnsureNotScalar(natStorageShape);
|
||||
OP_CHECK_IF(natShape.GetDimNum() != 1, OP_LOGE(context, "numAcceptedTokens must be 1D"),
|
||||
return ge::GRAPH_FAILED);
|
||||
OP_CHECK_IF(natShape.GetDim(0) != batch, OP_LOGE(context, "numAcceptedTokens.size must equal batch"),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
if (inputMode == 1) {
|
||||
const int64_t reqStateLen = (width - 1) + (seqLen - 1);
|
||||
OP_CHECK_IF(stateLen < reqStateLen,
|
||||
OP_LOGE(context,
|
||||
"spec decode requires stateLen >= (width-1) + (seqlen-1), got stateLen=%ld req=%ld",
|
||||
stateLen, reqStateLen),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
tiling.hasNumAcceptedTokens = 1;
|
||||
}
|
||||
}
|
||||
|
||||
tiling.hasBias = 0;
|
||||
auto biasShapePtr = context->GetOptionalInputShape(BIAS_INDEX);
|
||||
if (biasShapePtr != nullptr) {
|
||||
const auto biasStorageShape = biasShapePtr->GetStorageShape();
|
||||
const int64_t biasDimNum = biasStorageShape.GetDimNum();
|
||||
const bool biasAbsent = (biasDimNum == 0) || (biasDimNum == 1 && biasStorageShape.GetDim(0) <= 0);
|
||||
if (!biasAbsent) {
|
||||
auto biasShape = EnsureNotScalar(biasStorageShape);
|
||||
OP_CHECK_IF(biasShape.GetDimNum() != 1, OP_LOGE(context, "bias must be 1D: (dim,)"),
|
||||
return ge::GRAPH_FAILED);
|
||||
OP_CHECK_IF(biasShape.GetDim(0) != dim, OP_LOGE(context, "bias.size must equal dim"),
|
||||
return ge::GRAPH_FAILED);
|
||||
tiling.hasBias = 1;
|
||||
}
|
||||
}
|
||||
|
||||
const std::set<ge::DataType> supportedXDtype = {ge::DT_FLOAT16};
|
||||
auto xDesc = context->GetInputDesc(X_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, xDesc);
|
||||
const ge::DataType xDtype = xDesc->GetDataType();
|
||||
OP_CHECK_IF(supportedXDtype.count(xDtype) == 0, OP_LOGE(context, "x dtype only supports fp16"),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
auto wDesc = context->GetInputDesc(WEIGHT_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, wDesc);
|
||||
OP_CHECK_IF(wDesc->GetDataType() != xDtype, OP_LOGE(context, "weight dtype must equal x dtype"),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
if (tiling.hasBias == 1) {
|
||||
auto biasDesc = context->GetOptionalInputDesc(BIAS_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, biasDesc);
|
||||
OP_CHECK_IF(biasDesc->GetDataType() != xDtype, OP_LOGE(context, "bias dtype must equal x dtype"),
|
||||
return ge::GRAPH_FAILED);
|
||||
}
|
||||
|
||||
auto sDesc = context->GetInputDesc(CONV_STATES_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, sDesc);
|
||||
OP_CHECK_IF(sDesc->GetDataType() != xDtype, OP_LOGE(context, "convStates dtype must equal x dtype"),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
if (!qslAbsent) {
|
||||
auto qslDesc2 = context->GetOptionalInputDesc(QUERY_START_LOC_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, qslDesc2);
|
||||
const ge::DataType qslDtype = qslDesc2->GetDataType();
|
||||
OP_CHECK_IF(qslDtype != ge::DT_INT32 && qslDtype != ge::DT_INT64,
|
||||
OP_LOGE(context, "queryStartLoc dtype must be int32 or int64"), return ge::GRAPH_FAILED);
|
||||
tiling.queryStartLocUseInt64 = (qslDtype == ge::DT_INT64) ? 1 : 0;
|
||||
}
|
||||
|
||||
tiling.cacheIndicesUseInt64 = 0;
|
||||
if (tiling.hasCacheIndices == 1) {
|
||||
auto ciDesc = context->GetOptionalInputDesc(CACHE_INDICES_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, ciDesc);
|
||||
const ge::DataType ciDtype = ciDesc->GetDataType();
|
||||
OP_CHECK_IF(ciDtype != ge::DT_INT32 && ciDtype != ge::DT_INT64,
|
||||
OP_LOGE(context, "cacheIndices dtype must be int32 or int64"), return ge::GRAPH_FAILED);
|
||||
tiling.cacheIndicesUseInt64 = (ciDtype == ge::DT_INT64) ? 1 : 0;
|
||||
}
|
||||
|
||||
if (tiling.hasInitialStateMode == 1) {
|
||||
auto ismDesc = context->GetOptionalInputDesc(INITIAL_STATE_MODE_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, ismDesc);
|
||||
const ge::DataType ismDtype = ismDesc->GetDataType();
|
||||
OP_CHECK_IF(ismDtype != ge::DT_BOOL && ismDtype != ge::DT_INT32 && ismDtype != ge::DT_INT64,
|
||||
OP_LOGE(context, "initialStateMode dtype must be bool, int32 or int64"), return ge::GRAPH_FAILED);
|
||||
tiling.initialStateModeDtype =
|
||||
(ismDtype == ge::DT_INT64) ? 2 : ((ismDtype == ge::DT_INT32) ? 1 : 0);
|
||||
}
|
||||
|
||||
tiling.numAcceptedTokensUseInt64 = 0;
|
||||
if (tiling.hasNumAcceptedTokens == 1) {
|
||||
OP_CHECK_IF(width != 4, OP_LOGE(context, "numAcceptedTokens is only supported for width=4 currently"),
|
||||
return ge::GRAPH_FAILED);
|
||||
auto natDesc = context->GetOptionalInputDesc(NUM_ACCEPTED_TOKENS_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, natDesc);
|
||||
const ge::DataType natDtype = natDesc->GetDataType();
|
||||
OP_CHECK_IF(natDtype != ge::DT_INT32 && natDtype != ge::DT_INT64,
|
||||
OP_LOGE(context, "numAcceptedTokens dtype must be int32 or int64"), return ge::GRAPH_FAILED);
|
||||
tiling.numAcceptedTokensUseInt64 = (natDtype == ge::DT_INT64) ? 1 : 0;
|
||||
}
|
||||
|
||||
tiling.dim = dim;
|
||||
tiling.cuSeqlen = cuSeqlen;
|
||||
tiling.seqLen = seqLen;
|
||||
tiling.inputMode = inputMode;
|
||||
tiling.width = width;
|
||||
tiling.stateLen = stateLen;
|
||||
tiling.numCacheLines = numCacheLines;
|
||||
tiling.batch = batch;
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
static ge::graphStatus CausalConv1dTilingFunc(gert::TilingContext *context)
|
||||
{
|
||||
uint64_t ubSize;
|
||||
uint32_t coreNum;
|
||||
OP_CHECK_IF(GetPlatformInfo(context, ubSize, coreNum) != ge::GRAPH_SUCCESS,
|
||||
OP_LOGE(context, "GetPlatformInfo error"), return ge::GRAPH_FAILED);
|
||||
|
||||
OP_CHECK_IF(GetWorkspaceSize(context) != ge::GRAPH_SUCCESS, OP_LOGE(context, "GetWorkspaceSize error"),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
CausalConv1dTilingData *tiling = context->GetTilingData<CausalConv1dTilingData>();
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, tiling);
|
||||
OP_CHECK_IF(memset_s(tiling, sizeof(CausalConv1dTilingData), 0, sizeof(CausalConv1dTilingData)) != EOK,
|
||||
OP_LOGE(context, "set tiling data error"), return ge::GRAPH_FAILED);
|
||||
|
||||
OP_CHECK_IF(GetAttrsInfo(context, tiling->activationMode, tiling->padSlotId, tiling->runMode) != ge::GRAPH_SUCCESS,
|
||||
OP_LOGE(context, "GetAttrsInfo error"), return ge::GRAPH_FAILED);
|
||||
|
||||
OP_CHECK_IF(GetShapeDtypeInfo(context, *tiling) != ge::GRAPH_SUCCESS, OP_LOGE(context, "GetShapeDtypeInfo error"),
|
||||
return ge::GRAPH_FAILED);
|
||||
|
||||
const int64_t dim = tiling->dim;
|
||||
const int64_t batch = tiling->batch;
|
||||
OP_CHECK_IF(dim <= 0 || batch <= 0, OP_LOGE(context, "dim/batch must be positive"), return ge::GRAPH_FAILED);
|
||||
|
||||
const DimTileChoice choice = ChooseDimTileSize(context, batch, dim, coreNum);
|
||||
OP_CHECK_IF(choice.dimTileSize <= 0 || choice.blocksPerSeq <= 0 || choice.gridSize <= 0,
|
||||
OP_LOGE(context, "invalid dim_tile_size selection"), return ge::GRAPH_FAILED);
|
||||
|
||||
const uint32_t blockDim =
|
||||
(choice.gridSize < static_cast<int64_t>(coreNum)) ? static_cast<uint32_t>(choice.gridSize) : coreNum;
|
||||
|
||||
OP_LOGD(context,
|
||||
"Tiling result: batch[%ld], dim[%ld], dimTileSize[%ld], blocksPerSeq[%ld], gridSize[%ld], blockDim[%u], "
|
||||
"coreNum[%u].",
|
||||
batch, dim, choice.dimTileSize, choice.blocksPerSeq, choice.gridSize, blockDim, coreNum);
|
||||
|
||||
context->SetBlockDim(blockDim);
|
||||
tiling->dimTileSize = choice.dimTileSize;
|
||||
tiling->blocksPerSeq = choice.blocksPerSeq;
|
||||
|
||||
const uint64_t tilingKey = GET_TPL_TILING_KEY(CAUSAL_CONV1D_TPL_SCH_MODE_DEFAULT);
|
||||
context->SetTilingKey(tilingKey);
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
static ge::graphStatus TilingParseForCausalConv1d(gert::TilingParseContext *context)
|
||||
{
|
||||
OP_LOGD(context, "Enter TilingParseForCausalConv1d.");
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
IMPL_OP_OPTILING(CausalConv1dV310)
|
||||
.Tiling(CausalConv1dTilingFunc)
|
||||
.TilingParse<CausalConv1dCompileInfo>(TilingParseForCausalConv1d);
|
||||
} // namespace optiling
|
||||
@@ -0,0 +1,59 @@
|
||||
/**
|
||||
* 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.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file causal_conv1d_v310_tiling.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef ASCEND_OPS_CAUSAL_CONV1D_V310_TILING_DATA_H
|
||||
#define ASCEND_OPS_CAUSAL_CONV1D_V310_TILING_DATA_H
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
#include "register/tilingdata_base.h"
|
||||
#include "register/op_impl_registry.h"
|
||||
#include "tiling/platform/platform_ascendc.h"
|
||||
#include "platform/platform_infos_def.h"
|
||||
#include "tiling_base/error_log.h"
|
||||
|
||||
namespace optiling {
|
||||
|
||||
BEGIN_TILING_DATA_DEF(CausalConv1dTilingData)
|
||||
TILING_DATA_FIELD_DEF(int64_t, dim);
|
||||
TILING_DATA_FIELD_DEF(int64_t, cuSeqlen);
|
||||
TILING_DATA_FIELD_DEF(int64_t, seqLen);
|
||||
TILING_DATA_FIELD_DEF(int64_t, inputMode);
|
||||
|
||||
TILING_DATA_FIELD_DEF(int64_t, width);
|
||||
|
||||
TILING_DATA_FIELD_DEF(int64_t, stateLen);
|
||||
TILING_DATA_FIELD_DEF(int64_t, numCacheLines);
|
||||
|
||||
TILING_DATA_FIELD_DEF(int64_t, batch);
|
||||
|
||||
TILING_DATA_FIELD_DEF(int64_t, activationMode);
|
||||
TILING_DATA_FIELD_DEF(int64_t, padSlotId);
|
||||
|
||||
TILING_DATA_FIELD_DEF(int64_t, hasBias);
|
||||
|
||||
TILING_DATA_FIELD_DEF(int64_t, dimTileSize);
|
||||
TILING_DATA_FIELD_DEF(int64_t, blocksPerSeq);
|
||||
END_TILING_DATA_DEF;
|
||||
struct CausalConv1dCompileInfo {
|
||||
uint64_t ubSize = 0;
|
||||
uint32_t coreNum = 0;
|
||||
};
|
||||
REGISTER_TILING_DATA_CLASS(CausalConv1dV310, CausalConv1dTilingData)
|
||||
|
||||
} // namespace optiling
|
||||
|
||||
#endif // ASCEND_OPS_CAUSAL_CONV1D_V310_TILING_DATA_H
|
||||
61
csrc/moe/causal_conv1d_v310/op_host/math_util.h
Normal file
61
csrc/moe/causal_conv1d_v310/op_host/math_util.h
Normal file
@@ -0,0 +1,61 @@
|
||||
/**
|
||||
* 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 math_util.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef TILING_MATMUL_MATH_UTIL_H
|
||||
#define TILING_MATMUL_MATH_UTIL_H
|
||||
|
||||
#include <array>
|
||||
#include <cstdint>
|
||||
#include <vector>
|
||||
#include <utility>
|
||||
namespace matmul_tiling {
|
||||
class MathUtil {
|
||||
public:
|
||||
static bool IsEqual(float leftValue, float rightValue);
|
||||
template<typename T>
|
||||
static auto CeilDivision(T num1, T num2) -> T
|
||||
{
|
||||
if (num2 == 0) {
|
||||
return 0;
|
||||
}
|
||||
return static_cast<T>((static_cast<int64_t>(num1) + static_cast<int64_t>(num2) - 1) /
|
||||
static_cast<int64_t>(num2));
|
||||
}
|
||||
template<typename T>
|
||||
static auto Align(T num1, T num2) -> T
|
||||
{
|
||||
return CeilDivision(num1, num2) * num2;
|
||||
}
|
||||
static int32_t AlignDown(int32_t num1, int32_t num2);
|
||||
static bool CheckMulOverflow(int32_t a, int32_t b, int32_t &c);
|
||||
static int32_t MapShape(int32_t shape, bool roundUpFlag = true);
|
||||
static void AddFactor(std::vector<int32_t> &dimsFactors, int32_t dim);
|
||||
static void GetFactorCnt(const int32_t shape, int32_t &factorCnt, const int32_t factorStart,
|
||||
const int32_t factorEnd);
|
||||
static void GetFactorLayerCnt(const int32_t shape, int32_t &factorCnt, const int32_t factorStart,
|
||||
const int32_t factorEnd);
|
||||
static bool CheckFactorNumSatisfy(const int32_t dim);
|
||||
static int32_t FindBestSingleCore(const int32_t oriShape, const int32_t mappedShape, const int32_t coreNum,
|
||||
bool isKDim);
|
||||
static void GetFactors(std::vector<int32_t> &factorList, int32_t srcNum, int32_t minFactor, int32_t maxFactor);
|
||||
static void GetFactors(std::vector<int32_t> &factorList, int32_t srcNum, int32_t maxFactor);
|
||||
static void GetBlockFactors(std::vector<int32_t> &factorList, const int32_t oriShape, const int32_t mpShape,
|
||||
const int32_t coreNum, const int32_t maxNum);
|
||||
static int32_t GetNonFactorMap(std::vector<int32_t> &factorList, int32_t srcNum, int32_t maxFactor);
|
||||
static std::vector<std::pair<int, int>> GetFactorPairs(int32_t num);
|
||||
static std::pair<int32_t, int32_t> DivideIntoMainAndTail(int32_t num, int32_t divisor);
|
||||
};
|
||||
} // namespace matmul_tiling
|
||||
#endif // _MATH_UTIL_H_
|
||||
Reference in New Issue
Block a user