EngineX-Ascend/enginex-ascend-910-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

init v0.23.0

Browse Source 
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
Sun Ruoxi
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
csrc/moe/chunk_fwd_o/op_host/CMakeLists.txt Normal file
View File

@@ -0,0 +1,32 @@
set(CURRENT_CMAKE_DIR ${CMAKE_CURRENT_SOURCE_DIR})
set(CATLASS_INCLUDE_DIR "${CMAKE_SOURCE_DIR}/third_party/catlass/include")
get_filename_component(CATLASS_INCLUDE_DIR_ABS ${CATLASS_INCLUDE_DIR} ABSOLUTE)
set(COMMON_KERNEL_UTILS_DIR "${CMAKE_SOURCE_DIR}/moe/common")
get_filename_component(COMMON_KERNEL_UTILS_DIR_ABS ${COMMON_KERNEL_UTILS_DIR} ABSOLUTE)
add_op_to_compiled_list()
if (BUILD_OPEN_PROJECT)
target_sources(op_host_aclnnExc PRIVATE
chunk_fwd_o_def.cpp
)
endif()
add_ops_compile_options(
OP_NAME ChunkFwdO
OPTIONS
--cce-auto-sync=off
-Wno-deprecated-declarations
-I${CATLASS_INCLUDE_DIR_ABS}
-I${COMMON_KERNEL_UTILS_DIR_ABS}
)
if (NOT BUILD_OPS_RTY_KERNEL)
add_modules_sources(OPTYPE chunk_fwd_o ACLNNTYPE aclnn_exclude)
target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}
${CATLASS_INCLUDE_DIR_ABS}
)
endif()

105
csrc/moe/chunk_fwd_o/op_host/chunk_fwd_o_def.cpp Normal file
View File

@@ -0,0 +1,105 @@
/**
 * Copyright (c) 2025 Tianjin University, Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * the BSD 3-Clause License (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.
 */
/*!
*\file chunk_fwd_o_def.cpp
*\brief
*/
#include "register/op_def_registry.h"
namespace ops {
class ChunkFwdO : public OpDef {
public:
explicit ChunkFwdO(const char *name) : OpDef(name)
{
// Define inputs
this->Input("q")
.ParamType(REQUIRED)
.DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16})
.Format({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})
.AutoContiguous();
this->Input("k")
.ParamType(REQUIRED)
.DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16})
.Format({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})
.AutoContiguous();
this->Input("v")
.ParamType(REQUIRED)
.DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16})
.Format({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})
.AutoContiguous();
this->Input("h")
.ParamType(REQUIRED)
.DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16})
.Format({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})
.AutoContiguous();
this->Input("g")
.ParamType(REQUIRED)
.DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16})
.Format({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})
.AutoContiguous();
this->Input("cu_seqlens")
.ParamType(OPTIONAL)
.ValueDepend(OPTIONAL)
.DataType({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})
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
.AutoContiguous();
this->Input("chunk_offsets")
.ParamType(OPTIONAL)
.ValueDepend(OPTIONAL)
.DataType({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})
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
.AutoContiguous();
this->Output("o")
.ParamType(REQUIRED)
.DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16})
.Format({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});
this->Attr("scale").AttrType(REQUIRED).Float(1.0);
this->Attr("chunk_size").AttrType(REQUIRED).Int(64);
OpAICoreConfig aicore_config;
aicore_config.DynamicCompileStaticFlag(true)
.DynamicFormatFlag(true)
.DynamicRankSupportFlag(true)
.DynamicShapeSupportFlag(true)
.NeedCheckSupportFlag(false)
.PrecisionReduceFlag(true)
.ExtendCfgInfo("prebuildPattern.value", "Opaque")
.ExtendCfgInfo("coreType.value", "AiCore")
.ExtendCfgInfo("jitCompile.flag", "static_false,dynamic_false");
this->AICore().AddConfig("ascend910b", aicore_config);
this->AICore().AddConfig("ascend910_93", aicore_config);
this->AICore().AddConfig("ascend950", aicore_config);
this->AICore().AddConfig("ascend310p", aicore_config);
}
};
OP_ADD(ChunkFwdO);
} // namespace ops

157
csrc/moe/chunk_fwd_o/op_host/chunk_fwd_o_tiling.cpp Normal file
View File

@@ -0,0 +1,157 @@
/**
 * Copyright (c) 2025 Tianjin University, Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * the BSD 3-Clause License (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.
 */
/*!
* \file chunk_fwd_o_tiling.cpp
* \brief
*/
#include "chunk_fwd_o_tiling.h"
#include <register/op_impl_registry.h>
#include "../tiling_base/data_copy_transpose_tiling.h"
#include "../tiling_base/tiling_templates_registry.h"
namespace optiling {
static constexpr size_t INPUT_Q_IDX = 0;
static constexpr size_t INPUT_K_IDX = 1;
static constexpr size_t INPUT_V_IDX = 2;
static constexpr size_t INPUT_H_IDX = 3;
static constexpr size_t INPUT_G_IDX = 4;
static constexpr size_t INPUT_SEQLENS_IDX = 5;
static constexpr size_t INPUT_CHUNK_INDICES_IDX = 6;
static constexpr size_t ATTR_SCALE_IDX = 0;
static constexpr size_t ATTR_CHUNK_SIZE_IDX = 1;
static constexpr size_t DIM_BATCH = 0;
static constexpr size_t DIM_HEAD_NUM = 1;
static constexpr size_t DIM_SEQLEN = 2;
static constexpr size_t DIM_HEAD_DIM = 3;
static void ChunkFwdOTilingDataPrint(gert::TilingContext *context, ChunkFwdOTilingData &tiling)
{
auto nodeName = context->GetNodeName();
OP_LOGD(nodeName, ">>>>>>>>>>>>>>> Start to print ChunkFwdO tiling data <<<<<<<<<<<<<<<<");
OP_LOGD(nodeName, "=== batch: %ld", tiling.get_shapeBatch());
OP_LOGD(nodeName, "=== seqlen: %ld", tiling.get_seqlen());
OP_LOGD(nodeName, "=== kNumHead: %ld", tiling.get_kNumHead());
OP_LOGD(nodeName, "=== vNumHead: %ld", tiling.get_vNumHead());
OP_LOGD(nodeName, "=== kHeadDim: %ld", tiling.get_kHeadDim());
OP_LOGD(nodeName, "=== vHeadDim: %ld", tiling.get_vHeadDim());
OP_LOGD(nodeName, "=== chunkSize: %ld", tiling.get_chunkSize());
OP_LOGD(nodeName, "=== dataType: %ld", tiling.get_dataType());
OP_LOGD(nodeName, "=== isVariedLen: %ld", tiling.get_isVariedLen());
OP_LOGD(nodeName, "=== tokenBatch: %f", tiling.get_tokenBatch());
OP_LOGD(nodeName, ">>>>>>>>>>>>>>> Print ChunkFwdO tiling data end <<<<<<<<<<<<<<<<");
}
ge::graphStatus Tiling4ChunkFwdO(gert::TilingContext *context)
{
OP_LOGD(context->GetNodeName(), "Tiling4ChunkFwdO start.");
ChunkFwdOTilingData tiling;
gert::Shape qStorageShape = context->GetOptionalInputShape(INPUT_Q_IDX)->GetStorageShape();
gert::Shape vStorageShape = context->GetOptionalInputShape(INPUT_V_IDX)->GetStorageShape();
int64_t seqlen = qStorageShape.GetDim(DIM_SEQLEN);
int64_t kNumHead = qStorageShape.GetDim(DIM_HEAD_NUM);
int64_t vNumHead = vStorageShape.GetDim(DIM_HEAD_NUM);
int64_t kHeadDim = qStorageShape.GetDim(DIM_HEAD_DIM);
int64_t vHeadDim = vStorageShape.GetDim(DIM_HEAD_DIM);
int64_t isVariedLen, shapeBatch, tokenBatch;
auto cuSeqlensTensor = context->GetOptionalInputTensor(INPUT_SEQLENS_IDX);
if (cuSeqlensTensor == nullptr) {
isVariedLen = false;
shapeBatch = qStorageShape.GetDim(DIM_BATCH);
tokenBatch = 1;
} else {
isVariedLen = true;
shapeBatch = 1;
tokenBatch = cuSeqlensTensor->GetStorageShape().GetDim(DIM_BATCH) - 1;
}
auto attrPtr = context->GetAttrs();
float scale = *(attrPtr->GetAttrPointer<double>(ATTR_SCALE_IDX));
int64_t chunkSize = *(attrPtr->GetAttrPointer<int64_t>(ATTR_CHUNK_SIZE_IDX));
auto dtype = context->GetInputTensor(0)->GetDataType();
uint64_t dataType = dtype == ge::DT_BF16 ? 1 : 0;
auto gDType = context->GetOptionalInputTensor(INPUT_G_IDX)->GetDataType();
int64_t gDataType = 2;
if (gDType == ge::DT_BF16) {
gDataType = 1;
} else if (gDType == ge::DT_FLOAT16) {
gDataType = 0;
}
const auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
uint32_t aicCoreNum = ascendcPlatform.GetCoreNumAic();
context->SetBlockDim(aicCoreNum);
constexpr size_t WORKSPACE_RSV_BYTE = 16 * 1024 * 1024;
constexpr size_t GM_ALIGN = 512;
int64_t pingpongStages = 2;
size_t workspaceOffset = ascendcPlatform.GetLibApiWorkSpaceSize();
workspaceOffset += WORKSPACE_RSV_BYTE;
tiling.set_vWorkspaceOffset(workspaceOffset);
workspaceOffset += (aicCoreNum * chunkSize * vHeadDim * sizeof(float) * pingpongStages + GM_ALIGN) / GM_ALIGN * GM_ALIGN;
tiling.set_hWorkspaceOffset(workspaceOffset);
workspaceOffset += (aicCoreNum * chunkSize * vHeadDim * sizeof(float) * pingpongStages + GM_ALIGN) / GM_ALIGN * GM_ALIGN;
tiling.set_attnWorkspaceOffset(workspaceOffset);
workspaceOffset += (aicCoreNum * chunkSize * chunkSize * sizeof(float) * pingpongStages + GM_ALIGN) / GM_ALIGN * GM_ALIGN;
tiling.set_aftermaskWorkspaceOffset(workspaceOffset);
workspaceOffset += (aicCoreNum * chunkSize * chunkSize * sizeof(float) * pingpongStages + GM_ALIGN) / GM_ALIGN * GM_ALIGN;
tiling.set_maskWorkspaceOffset(workspaceOffset);
workspaceOffset += (chunkSize * chunkSize + GM_ALIGN) / GM_ALIGN * GM_ALIGN;
workspaceOffset += WORKSPACE_RSV_BYTE;
size_t *currentWorkspace = context->GetWorkspaceSizes(1);
currentWorkspace[0] = (workspaceOffset - 0);
tiling.set_shapeBatch(shapeBatch);
tiling.set_seqlen(seqlen);
tiling.set_kNumHead(kNumHead);
tiling.set_vNumHead(vNumHead);
tiling.set_kHeadDim(kHeadDim);
tiling.set_vHeadDim(vHeadDim);
tiling.set_scale(scale);
tiling.set_chunkSize(chunkSize);
tiling.set_isVariedLen(isVariedLen);
tiling.set_tokenBatch(tokenBatch);
tiling.set_dataType(dataType);
tiling.set_gDataType(gDataType);
tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
ChunkFwdOTilingDataPrint(context, tiling);
OP_LOGD(context->GetNodeName(), "Tiling4ChunkFwdO end.");
return ge::GRAPH_SUCCESS;
}
ge::graphStatus TilingPrepareForChunkFwdO(gert::TilingParseContext *context)
{
(void)context;
return ge::GRAPH_SUCCESS;
}
IMPL_OP_OPTILING(ChunkFwdO)
.Tiling(Tiling4ChunkFwdO)
.TilingParse<ChunkFwdOCompileInfo>(TilingPrepareForChunkFwdO);
} // namespace optiling

46
csrc/moe/chunk_fwd_o/op_host/chunk_fwd_o_tiling.h Normal file
View File

@@ -0,0 +1,46 @@
/**
 * Copyright (c) 2025 Tianjin University, Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * the BSD 3-Clause License (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.
 */
/*!
* \file chunk_fwd_o_tiling.h
* \brief
*/
#pragma once
#include <cstdint>
#include <register/tilingdata_base.h>
#include <tiling/tiling_api.h>
namespace optiling {
BEGIN_TILING_DATA_DEF(ChunkFwdOTilingData)
TILING_DATA_FIELD_DEF(int64_t, shapeBatch);
TILING_DATA_FIELD_DEF(int64_t, seqlen);
TILING_DATA_FIELD_DEF(int64_t, kNumHead);
TILING_DATA_FIELD_DEF(int64_t, vNumHead);
TILING_DATA_FIELD_DEF(int64_t, kHeadDim);
TILING_DATA_FIELD_DEF(int64_t, vHeadDim);
TILING_DATA_FIELD_DEF(int64_t, chunkSize);
TILING_DATA_FIELD_DEF(int64_t, isVariedLen);
TILING_DATA_FIELD_DEF(int64_t, tokenBatch);
TILING_DATA_FIELD_DEF(int64_t, dataType);
TILING_DATA_FIELD_DEF(int64_t, gDataType);
TILING_DATA_FIELD_DEF(int64_t, vWorkspaceOffset);
TILING_DATA_FIELD_DEF(int64_t, hWorkspaceOffset);
TILING_DATA_FIELD_DEF(int64_t, attnWorkspaceOffset);
TILING_DATA_FIELD_DEF(int64_t, aftermaskWorkspaceOffset);
TILING_DATA_FIELD_DEF(int64_t, maskWorkspaceOffset);
TILING_DATA_FIELD_DEF(float, scale);
END_TILING_DATA_DEF;
REGISTER_TILING_DATA_CLASS(ChunkFwdO, ChunkFwdOTilingData)
struct ChunkFwdOCompileInfo {};
} // namespace optiling

162
csrc/moe/chunk_fwd_o/op_host/op_api/aclnn_chunk_fwd_o.cpp Normal file
View File

@@ -0,0 +1,162 @@
/**
* Copyright (c) 2025 Tianjin University, 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.
*/
#include "aclnn_chunk_fwd_o.h"
#include "chunk_fwd_o.h"
#include <dlfcn.h>
#include <new>
#include "aclnn_kernels/transdata.h"
#include "aclnn_kernels/contiguous.h"
#include "acl/acl.h"
#include "aclnn/aclnn_base.h"
#include "aclnn_kernels/common/op_error_check.h"
#include "opdev/common_types.h"
#include "opdev/data_type_utils.h"
#include "opdev/format_utils.h"
#include "opdev/op_dfx.h"
#include "opdev/op_executor.h"
#include "opdev/op_log.h"
#include "opdev/platform.h"
#include "opdev/shape_utils.h"
#include "opdev/tensor_view_utils.h"
#include "opdev/make_op_executor.h"
using namespace op;
#ifdef __cplusplus
extern "C" {
#endif
struct ChunkFwdOParams {
const aclTensor *q = nullptr;
const aclTensor *k = nullptr;
const aclTensor *v = nullptr;
const aclTensor *h = nullptr;
const aclTensor *g = nullptr;
const aclIntArray *cuSeqlensOptional = nullptr;
const aclIntArray *chunkOffsetsOptional = nullptr;
double scale = 1.0;
int64_t chunkSize = 64;
const aclTensor *oOut = nullptr;
};
static aclnnStatus CheckNotNull(ChunkFwdOParams params)
{
CHECK_COND(params.q != nullptr, ACLNN_ERR_PARAM_NULLPTR, "q must not be nullptr.");
CHECK_COND(params.k != nullptr, ACLNN_ERR_PARAM_NULLPTR, "k must not be nullptr.");
CHECK_COND(params.v != nullptr, ACLNN_ERR_PARAM_NULLPTR, "v must not be nullptr.");
CHECK_COND(params.h != nullptr, ACLNN_ERR_PARAM_NULLPTR, "h must not be nullptr.");
CHECK_COND(params.g != nullptr, ACLNN_ERR_PARAM_NULLPTR, "g must not be nullptr.");
CHECK_COND(params.oOut != nullptr, ACLNN_ERR_PARAM_NULLPTR, "oOut must not be nullptr.");
return ACLNN_SUCCESS;
}
static aclnnStatus CheckFormat(ChunkFwdOParams params)
{
return ACLNN_SUCCESS;
}
static aclnnStatus CheckShape(ChunkFwdOParams params)
{
return ACLNN_SUCCESS;
}
static aclnnStatus DataContiguous(const aclTensor *&tensor, aclOpExecutor *executor)
{
tensor = l0op::Contiguous(tensor, executor);
CHECK_RET(tensor != nullptr, ACLNN_ERR_INNER_NULLPTR);
return ACLNN_SUCCESS;
}
static aclnnStatus ParamsDataContiguous(ChunkFwdOParams &params, aclOpExecutor *executorPtr)
{
CHECK_COND(DataContiguous(params.q, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID,
"Contiguous q failed.");
CHECK_COND(DataContiguous(params.k, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID,
"Contiguous k failed.");
CHECK_COND(DataContiguous(params.v, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID,
"Contiguous v failed.");
CHECK_COND(DataContiguous(params.h, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID,
"Contiguous h failed.");
CHECK_COND(DataContiguous(params.g, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID,
"Contiguous g failed.");
return ACLNN_SUCCESS;
}
static aclnnStatus CheckDtype(ChunkFwdOParams params)
{
return ACLNN_SUCCESS;
}
static aclnnStatus CheckParams(ChunkFwdOParams params)
{
CHECK_RET(CheckNotNull(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
CHECK_RET(CheckFormat(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
CHECK_RET(CheckShape(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
CHECK_RET(CheckDtype(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
return ACLNN_SUCCESS;
}
aclnnStatus aclnnChunkFwdOGetWorkspaceSize(
const aclTensor *q,
const aclTensor *k,
const aclTensor *v,
const aclTensor *h,
const aclTensor *g,
const aclIntArray *cuSeqlensOptional,
const aclIntArray *chunkOffsetsOptional,
double scale,
int64_t chunkSize,
const aclTensor *oOut,
uint64_t *workspaceSize,
aclOpExecutor **executor)
{
ChunkFwdOParams params{q, k, v, h, g, cuSeqlensOptional, chunkOffsetsOptional, scale, chunkSize, oOut};
// Standard syntax, Check parameters.
L2_DFX_PHASE_1(aclnnChunkFwdO, DFX_IN(q, k, v, h, g, cuSeqlensOptional, chunkOffsetsOptional),
DFX_OUT(oOut));
// 固定写法，创建OpExecutor
auto uniqueExecutor = CREATE_EXECUTOR();
CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);
auto executorPtr = uniqueExecutor.get();
// 固定写法，参数检查
auto ret = CheckParams(params);
CHECK_RET(ret == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
CHECK_COND(ParamsDataContiguous(params, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID,
"ParamsDataContiguous failed.");
auto result = l0op::ChunkFwdO(params.q, params.k, params.v, params.h, params.g, params.cuSeqlensOptional, params.chunkOffsetsOptional, params.scale, params.chunkSize, params.oOut, executorPtr);
CHECK_RET(result[0] != nullptr, ACLNN_ERR_PARAM_NULLPTR);
// If the output tensor is non-contiguous, convert the calculated contiguous tensor to non-contiguous.
auto viewCopyResult = l0op::ViewCopy(result[0], params.oOut, executorPtr);
CHECK_RET(viewCopyResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
// Standard syntax, get the size of workspace needed during computation.
*workspaceSize = uniqueExecutor->GetWorkspaceSize();
uniqueExecutor.ReleaseTo(executor);
return ACLNN_SUCCESS;
}
aclnnStatus aclnnChunkFwdO(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)
{
L2_DFX_PHASE_2(aclnnChunkFwdO);
CHECK_COND(CommonOpExecutorRun(workspace, workspaceSize, executor, stream) == ACLNN_SUCCESS, ACLNN_ERR_INNER,
"This is an error in ChunkFwdO launch aicore.");
return ACLNN_SUCCESS;
}
#ifdef __cplusplus
}
#endif

65
csrc/moe/chunk_fwd_o/op_host/op_api/aclnn_chunk_fwd_o.h Normal file
View File

@@ -0,0 +1,65 @@
/**
* Copyright (c) 2025 Tianjin University, 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.
*/
#ifndef OP_API_INC_ACLNN_CHUNK_FWD_O_H
#define OP_API_INC_ACLNN_CHUNK_FWD_O_H
#include "aclnn/aclnn_base.h"
#ifdef __cplusplus
extern "C" {
#endif
/* function: aclnnChunkFwdOGetWorkspaceSize
* parameters :
* q : required
* k : required
* v : required
* h : required
* g : required
* cuSeqlensOptional : optional
* chunkOffsetsOptional : optional
* scale : required
* chunkSize : required
* oOut : required
* workspaceSize : size of workspace(output).
* executor : executor context(output).
*/
__attribute__((visibility("default")))
aclnnStatus aclnnChunkFwdOGetWorkspaceSize(
const aclTensor *q,
const aclTensor *k,
const aclTensor *v,
const aclTensor *h,
const aclTensor *g,
const aclIntArray *cuSeqlensOptional,
const aclIntArray *chunkOffsetsOptional,
double scale,
int64_t chunkSize,
const aclTensor *oOut,
uint64_t *workspaceSize,
aclOpExecutor **executor);
/* function: aclnnChunkFwdO
* parameters :
* workspace : workspace memory addr(input).
* workspaceSize : size of workspace(input).
* executor : executor context(input).
* stream : acl stream.
*/
__attribute__((visibility("default")))
aclnnStatus aclnnChunkFwdO(
void *workspace,
uint64_t workspaceSize,
aclOpExecutor *executor,
aclrtStream stream);
#ifdef __cplusplus
}
#endif
#endif

67
csrc/moe/chunk_fwd_o/op_host/op_api/chunk_fwd_o.cpp Normal file
View File

@@ -0,0 +1,67 @@
/**
* Copyright (c) 2025 Tianjin University, 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.
*/
#include "opdev/op_log.h"
#include "opdev/op_dfx.h"
#include "opdev/make_op_executor.h"
#include "chunk_fwd_o.h"
using namespace op;
namespace l0op {
OP_TYPE_REGISTER(ChunkFwdO);
const std::array<const aclTensor *, 1> ChunkFwdO(
const aclTensor *q,
const aclTensor *k,
const aclTensor *v,
const aclTensor *h,
const aclTensor *g,
const aclIntArray *cuSeqlensOptional,
const aclIntArray *chunkOffsetsOptional,
double scale,
int64_t chunkSize,
const aclTensor *oOut,
aclOpExecutor *executor)
{
L0_DFX(ChunkFwdO, q, k, v, h, g, cuSeqlensOptional, chunkOffsetsOptional, scale, chunkSize, oOut);
const aclTensor *actualCuSeqlens = nullptr;
if (cuSeqlensOptional) {
actualCuSeqlens = executor->ConvertToTensor(cuSeqlensOptional, DataType::DT_INT64);
const_cast<aclTensor *>(actualCuSeqlens)->SetStorageFormat(Format::FORMAT_ND);
const_cast<aclTensor *>(actualCuSeqlens)->SetViewFormat(Format::FORMAT_ND);
const_cast<aclTensor *>(actualCuSeqlens)->SetOriginalFormat(Format::FORMAT_ND);
} else {
actualCuSeqlens = nullptr;
}
const aclTensor *actualChunkOffsets = nullptr;
if (chunkOffsetsOptional) {
actualChunkOffsets = executor->ConvertToTensor(chunkOffsetsOptional, DataType::DT_INT64);
const_cast<aclTensor *>(actualChunkOffsets)->SetStorageFormat(Format::FORMAT_ND);
const_cast<aclTensor *>(actualChunkOffsets)->SetViewFormat(Format::FORMAT_ND);
const_cast<aclTensor *>(actualChunkOffsets)->SetOriginalFormat(Format::FORMAT_ND);
} else {
actualChunkOffsets = nullptr;
}
auto ret = ADD_TO_LAUNCHER_LIST_AICORE(ChunkFwdO,
OP_INPUT(q, k, v, h, g, actualCuSeqlens, actualChunkOffsets),
OP_OUTPUT(oOut),
OP_ATTR(scale, chunkSize));
if (ret != ACLNN_SUCCESS) {
OP_LOGE(ACLNN_ERR_PARAM_INVALID, "ADD_TO_LAUNCHER_LIST_AICORE failed.");
return {nullptr};
}
return {oOut};
}
} // namespace l0op

30
csrc/moe/chunk_fwd_o/op_host/op_api/chunk_fwd_o.h Normal file
View File

@@ -0,0 +1,30 @@
/**
* Copyright (c) 2025 Tianjin University, 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.
*/
#ifndef OP_API_INC_LEVEL0_OP_CHUNK_FWD_O_H
#define OP_API_INC_LEVEL0_OP_CHUNK_FWD_O_H
#include "opdev/op_executor.h"
namespace l0op {
const std::array<const aclTensor *, 1> ChunkFwdO(
const aclTensor *q,
const aclTensor *k,
const aclTensor *v,
const aclTensor *h,
const aclTensor *g,
const aclIntArray *cuSeqlensOptional,
const aclIntArray *chunkOffsetsOptional,
double scale,
int64_t chunkSize,
const aclTensor *oOut,
aclOpExecutor *executor);
}
#endif