在 CANN ops-transformer 中理解 DenseLightningIndexerKLLossGradMetadata用于负载均衡的分核 metadata 生成器【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读DenseLightningIndexerKLLossGradMetadata 是 CANN ops-transformer 中专门为aclnnDenseLightningIndexerKLLossGrad算子准备的前置 AICPU 算子。它接收主算子的 shape、layout、mask 与压缩比例信息基于每个 query 的有效序列长度估算负载将 B/S1 合轴后的任务均衡切分到可用 AIC core 上并输出一份固定布局的 INT32 metadata供主算子在 tiling 阶段直接消费从而减少对 host array 的访问。本文围绕该算子的官方文档展开结合仓库内的源码实现、示例与配套算子系统讲解它的定位、接口、参数约束、metadata 布局及调用方式帮助你理解如何在自己的模型中正确接入这条metadata 前置生成的流水线。产品支持情况仅在 Ascend 950 系列上可用该算子在当前仓库中明确标注为仅支持 Ascend 950PR/Ascend 950DT 系列产品文档表格与 README 保持一致其余产品线均不支持产品是否支持Ascend 950PR/Ascend 950DT支持Atlas A3 训练系列产品/Atlas A3 推理系列产品不支持Atlas A2 训练系列产品/Atlas A2 推理系列产品不支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持这与源码中的架构分派逻辑一致。在 dense_lightning_indexer_kl_loss_grad_metadata_aicpu.cpp 中内核首先从上下文中读取soc_version属性再分发到 arch35A5/Ascend 950 系实现而 Host 侧的参数校验 dense_lightning_indexer_kl_loss_grad_metadata_check.h 则先判断socVersion是否包含Ascend950仅在此前提下才执行完整的 A5 参数校验否则只做最基础的空指针检查并直接返回——从源码结构看该算子的完整能力仅面向 Ascend 950 系列。功能定位主算子的前置分核切分器该算子是一个 AICPU 算子接口功能定义如下根据aclnnDenseLightningIndexerKLLossGrad算子的输入 shape、layout、mask 和压缩比例信息计算并输出分核切分 metadata。输出结果可作为aclnnDenseLightningIndexerKLLossGrad算子的metadataOptional输入减少主算子 tiling 阶段对 host array 的访问。该算子不建议单独使用建议与aclnnDenseLightningIndexerKLLossGrad算子配合使用形成完整工作流接收主算子的 shape 信息包括batchSize、maxSeqLenQ、maxSeqLenK、numHeadsQ、numHeadsK、headDim、layout和mask信息。根据每个 query 对应的有效序列长度估算负载并将 B/S1 合轴后的任务均衡切分到可用 AIC core 上。输出 metadata 后后续作为aclnnDenseLightningIndexerKLLossGrad算子的metadataOptional输入使用。从 AICPU 内核实现 dense_lightning_indexer_kl_loss_grad_metadata_aicpu_arch35.h 可以看到核心计算逻辑CalcSplitInfo()遍历每个 batch累加GetS1SeqSize(bIdx)每个 batch 中 query 的有效序列长度得到合轴后的任务总行数totalNum。GenMetadata()按formerCoreProcessNum CeilDiv(totalNum, aicCoreNum)计算上对齐分核粒度remainCoreProcessNum formerCoreProcessNum - 1为下对齐分核粒度remainCoreNum为处理残余粒度的 core 数量usedCoreNum为实际使用的 AIC core 数量最终写入 metadata 的前 5 个字段。序列长度的获取优先级为seqused_q/seqused_k提供的实际参与运算长度 TND 场景下cu_seqlens的差分 标量max_seqlen见 GetS1SeqSize/GetS2SeqSize。内核注册信息位于 dense_lightning_indexer_kl_loss_grad_metadata_aicpu.json该算子运行在DNN_VM_AICPU引擎kernelSo为libtransformer_aicpu_kernels.socomputeCost与workspaceSize均配置为 100属于用户自定义userDefined: True的 AICPU kernel。函数原型与两段式调用每个算子分为两段式接口必须先调用aclnnDenseLightningIndexerKLLossGradMetadataGetWorkspaceSize获取 workspace 大小再调用aclnnDenseLightningIndexerKLLossGradMetadata执行计算。第一段接口原型aclnnStatus aclnnDenseLightningIndexerKLLossGradMetadataGetWorkspaceSize( const aclTensor *cuSeqLensQOptional, const aclTensor *cuSeqLensKOptional, const aclTensor *seqUsedQOptional, const aclTensor *seqUsedKOptional, const aclTensor *cmpResidualKOptional, int64_t batchSize, int64_t maxSeqLenQ, int64_t maxSeqLenK, int64_t numHeadsQ, int64_t numHeadsK, int64_t headDim, char *layoutQ, char *layoutK, int64_t maskMode, int64_t cmpRatio, const aclTensor *metadata, uint64_t *workspaceSize, aclOpExecutor **executor)第二段接口原型aclnnStatus aclnnDenseLightningIndexerKLLossGradMetadata( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)公开 API 的声明位于 aclnn_dense_lightning_indexer_kl_loss_grad_metadata.h宿主侧实现位于 aclnn_dense_lightning_indexer_kl_loss_grad_metadata.cpp第一段接口先对workspaceSize/executor做空指针检查再从当前平台信息中读取aicCoreNum、aivCoreNum与socVersion执行参数校验对 5 个可选 Tensor 入参分别做 Contiguous 处理最后构造执行器并返回 workspace 大小第二段接口通过CommonOpExecutorRun启动异步执行。第一段接口参数说明输入参数前 15 个参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 TensorcuSeqLensQOptionalconst aclTensor*输入表示不同 batch 中 query 的累积 sequence length。支持空 TensorTND 场景下必传并可通过该入参 shape 推导 batch第一个值固定为 0shape 固定为 (B1, )。INT32ND1 维(B1,)×cuSeqLensKOptionalconst aclTensor*输入表示不同 batch 中 key 的累积 sequence length。支持空 TensorTND 场景下必传第一个值固定为 0shape 固定为 (B1, )。INT32ND1 维(B1,)×seqUsedQOptionalconst aclTensor*输入表示不同 batch 中 query 实际参与运算的 sequence length。支持空 Tensorshape 固定为 (B, )。INT32ND1 维(B,)×seqUsedKOptionalconst aclTensor*输入表示不同 batch 中 key 实际参与运算的 sequence length。支持空 Tensorshape 固定为 (B, )。INT32ND1 维(B,)×cmpResidualKOptionalconst aclTensor*输入表示不同 batch 中 key 的 sequence length 与 cmpRatio 相关的残差。支持空 Tensorshape 固定为 (B, )。INT32ND1 维(B,)×batchSizeint64_t输入表示 batch 数量。支持非负数TND 场景可填 0并通过 cuSeqLensQOptional 推导建议值为 0。----maxSeqLenQint64_t输入表示 query 的最大 sequence length。支持非负数BSND 场景必须为正数建议值为 0。----maxSeqLenKint64_t输入表示 key 的最大 sequence length。支持非负数BSND 场景必须为正数建议值为 0。----numHeadsQint64_t输入表示 query 的 head 个数。必须为正数并且能被 numHeadsK 整除当前支持 [1, 128]。----numHeadsKint64_t输入表示 key 的 head 个数。必须为正数当前仅支持 1。----headDimint64_t输入表示 q/k 的 head dimension。必须为正数当前仅支持 128。----layoutQchar*输入表示 query 侧的排列格式。支持 BSND、TND建议值为 BSND。----layoutKchar*输入表示 key 侧的排列格式。支持 BSND、TND建议值为 BSND。----maskModeint64_t输入表示 mask 模式。0: No mask3: rightDownCausal 模式的 mask对应以右顶点为划分的下三角场景建议值为 0。----cmpRatioint64_t输入表示 key 的压缩率。取值范围 [1, 128]建议值 1表示无压缩。----Host 侧校验常量与文档给出的取值范围一一对应在 dense_lightning_indexer_kl_loss_grad_metadata_check.h 中定义了DLI_NO_MASK_MODE 0、DLI_CAUSAL_MASK_MODE 3、DLI_CMP_RATIO_LOWER_BOUND 1、DLI_CMP_RATIO_UPPER_BOUND 128、DLI_NUM_HEADS_Q_LOWER_BOUND_A5 1、DLI_NUM_HEADS_Q_UPPER_BOUND_A5 128以及DLIKG_METADATA_SIZE 64并在CheckSingleParamDli中对numHeadsK 1、headDim 128、layoutQ/layoutK ∈ {BSND, TND}等约束逐一强制校验。输出参数参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 Tensormetadataconst aclTensor*输出表示负载均衡结果输出。输出结果作为 aclnnDenseLightningIndexerKLLossGrad 的 metadataOptional 输入。INT32ND1 维shape 固定为 (64,)×workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程。-----metadata 的 shape 由 infershape 阶段直接写死为 64见 dense_lightning_indexer_kl_loss_grad_metadata_infershape.cpp其中InferShape将输出第 0 维设为optiling::DLIKG_METADATA_SIZE64InferDtype将输出类型固定为DT_INT32。返回值与错误码aclnnStatus返回状态码具体参见 aclnn 返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_INNER_NULLPTR561103workspaceSize 或 executor 为空指针。ACLNN_ERR_PARAM_INVALID多场景161002cuSeqLensQOptional、cuSeqLensKOptional、seqUsedQOptional、seqUsedKOptional 或 cmpResidualKOptional 的数据类型或维度不在支持范围内metadata 为空指针或其数据类型、维度或 shape 不在支持范围内numHeadsQ、numHeadsK 或 headDim 不在支持范围内layoutQ 或 layoutK 为空指针或不在支持范围内batchSize、maxSeqLenQ、maxSeqLenK、maskMode 或 cmpRatio 不满足取值约束TND 场景下未传入对应的 cuSeqLensQOptional 或 cuSeqLensKOptional或 maskMode3 且 cmpRatio 不为 1 时未传入 cmpResidualKOptional从 query、key 或 cmpResidualKOptional 推导出的 batch 数量不一致。从源码看第一段接口的空指针检查在 aclnn_dense_lightning_indexer_kl_loss_grad_metadata.cpp 中实现workspaceSize/executor为空即返回ACLNN_ERR_INNER_NULLPTR而参数范围与一致性校验则由ParamsCheckDli → ParamsCheckDliA5依次完成CheckSingleParamDli单参数取值、CheckExistenceDliTND 必传项与 mask/cmpRatio 联动必传项、CheckConsistencyDli维度、数据类型、batch 一致性三层校验。第二段接口参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnDenseLightningIndexerKLLossGradMetadataGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。返回值同为aclnnStatus具体参见 aclnn 返回码。约束说明确定性说明aclnnDenseLightningIndexerKLLossGradMetadata默认确定性实现。BSND 场景必传batchSize、maxSeqLenQ和maxSeqLenK参数以获取 shape 信息。TND 场景必传cuSeqLensQOptional、cuSeqLensKOptional参数以获取正确 shape 信息当batchSize为 0 时通过cuSeqLensQOptional的 shape 推导 batch。Maskmask_modemask_mode含义备注0无 mask。支持3rightDownCausal 模式的 mask对应以右顶点为划分的下三角场景。支持mask 模式的详细介绍可参考 sparse_mode 参数说明。特殊约束Batch 取值规则如果batchSize大于 0优先使用batchSize。如果batchSize小于等于 0且layoutQ为 TND则通过cuSeqLensQOptional的 shape 推导 batch。如果batchSize小于等于 0且layoutQ为 BSND则报错。Seqlen 取值规则TND 场景下通过cuSeqLensQOptional和cuSeqLensKOptional计算每个 batch 的实际 q/k 长度。BSND 场景下通过maxSeqLenQ和maxSeqLenK获取 q/k 长度。layout 约束layoutQ必须为 BSND 或 TND。layoutK支持 BSND 和 TND建议与layoutQ保持一致。head 约束numHeadsQ、numHeadsK和headDim必须为正数。numHeadsQ必须能被numHeadsK整除。dense 约束cmpRatio取值范围为 [1, 128]。maskMode当前仅支持 0 和 3。上述规则在 AICPU 内核的ParamsInit见 arch35 实现中有对应处理groupSize_ numHeadsQ_ / numHeadsK_计算 group sizemaskMode为 3RIGHT_DOWN_CAUSAL时attentionMode_置 1S2 基准粒度s2BaseSize_在 Ascend 910 系为 2048、Ascend 950 系为 128进一步印证了该算子在 950 系列上的专用适配。Metadata 输出布局metadata 输出为 INT32 Tensor当前 shape 固定为 (64,)字段布局如下字段index说明totalSize0B/S1 合轴后的任务总行数。formerCoreSize1上对齐场景下的分核粒度。remainCoreSize2下对齐场景下的分核粒度。remainCoreNum3处理 remainCoreSize 的 cube core 数量。usedCoreNum4主 kernel 实际使用的 cube core 数量。maxSeqK5key 的最大 sequence length。reserved6-63预留字段当前置 0。该布局与主算子侧共享的元数据定义完全一致在 dense_lightning_indexer_kl_loss_grad_metadata_arch35.h 中optiling命名空间定义了DLIKG_METADATA_SIZE 64、TOTAL_NUM 0、FORMER_CORE_PROCESS_NUM 1、REMAIN_CORE_PROCESS_NUM 2、REMAIN_CORE_NUM 3、USED_CORE_NUM 4等常量AICPU 内核正是按这些索引把负载切分结果写入gradMetadata[]并通过static_assert保证 64 个 INT32 足以承载DlikgMetadata结构。metadata 由前置算子生成、主算子消费前后两个算子通过这份固定协议完成分核切分信息的传递。调用示例以下代码演示了两段式完整调用流程TND 场景示例参数batchSize1、maxSeqLenQ16、maxSeqLenK4、numHeadsQ8、numHeadsK1、headDim128、maskMode0、cmpRatio4具体编译和执行过程请参见编译与运行样例#include algorithm #include cstdint #include cstdio #include cstring #include functional #include iostream #include utility #include vector #include acl/acl.h #include aclnnop/aclnn_dense_lightning_indexer_kl_loss_grad_metadata.h #define CHECK_LOG_RET(cond, ret_val, fmt, ...) \ do { \ if (!(cond)) { \ printf(fmt \n, ##__VA_ARGS__); \ return (ret_val); \ } \ } while (0) constexpr uint32_t DLI_METADATA_MAX_CORE_NUM 25; constexpr uint32_t DLI_METADATA_HEADER_SIZE 8; constexpr uint32_t DLI_METADATA_SIZE 64; struct DliGradKLLossMetaData { int32_t coreNum; int32_t totalSize; int32_t splitFactorSize; int32_t reserved[DLI_METADATA_HEADER_SIZE - 3]; int32_t bS1Index[DLI_METADATA_MAX_CORE_NUM]; }; struct ScopeGuard { explicit ScopeGuard(std::functionvoid() onExitScope) : m_exitFunc(std::move(onExitScope)), m_isDismissed(false) {} ScopeGuard(const ScopeGuard) delete; ScopeGuard operator(const ScopeGuard) delete; ~ScopeGuard() { if (!m_isDismissed) { m_exitFunc(); } } void Dismiss() { m_isDismissed true; } std::functionvoid() m_exitFunc; bool m_isDismissed; }; struct Tensor { void *hostAddr { nullptr }; void *deviceAddr { nullptr }; aclTensor *data { nullptr }; }; struct ArgScenario { bool hasCuSeq { true }; }; struct ArgContext { Tensor cuSeqLensQOptional {}; Tensor cuSeqLensKOptional {}; Tensor seqUsedQOptional {}; Tensor seqUsedKOptional {}; Tensor cmpResidualKOptional {}; Tensor metadata {}; int64_t batchSize { 0 }; int64_t maxSeqLenQ { 0 }; int64_t maxSeqLenK { 0 }; int64_t numHeadsQ { 8 }; int64_t numHeadsK { 1 }; int64_t headDim { 128 }; char *layoutQ { nullptr }; char *layoutK { nullptr }; int64_t maskMode { 0 }; int64_t cmpRatio { 4 }; }; int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } aclnnStatus Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclInit failed. ERROR: %d, ret); ret aclrtSetDevice(deviceId); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclrtSetDevice failed. ERROR: %d, ret); ret aclrtCreateStream(stream); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclrtCreateStream failed. ERROR: %d, ret); return ACL_SUCCESS; } void Finalize(int32_t deviceId, aclrtStream stream) { aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); } aclnnStatus CreateTensor(aclDataType dataType, const std::vectorint64_t shape, Tensor tensor) { auto size GetShapeSize(shape) * aclDataTypeSize(dataType); auto ret aclrtMallocHost((tensor.hostAddr), size); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclrtMallocHost failed. ERROR: %d, ret); memset(tensor.hostAddr, 0, size); ret aclrtMalloc((tensor.deviceAddr), size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclrtMalloc failed. ERROR: %d, ret); tensor.data aclCreateTensor(shape.data(), shape.size(), dataType, nullptr, 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), tensor.deviceAddr); ret aclrtMemcpy(tensor.deviceAddr, size, tensor.hostAddr, size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclrtMemcpy failed. ERROR: %d, ret); return ACL_SUCCESS; } void SetInt32TensorData(Tensor tensor, const std::vectorint32_t hostData) { auto size hostData.size() * sizeof(int32_t); memcpy(tensor.hostAddr, hostData.data(), size); aclrtMemcpy(tensor.deviceAddr, size, tensor.hostAddr, size, ACL_MEMCPY_HOST_TO_DEVICE); } void DestroyTensor(Tensor tensor) { if (tensor.data ! nullptr) { aclDestroyTensor(tensor.data); tensor.data nullptr; } if (tensor.deviceAddr ! nullptr) { aclrtFree(tensor.deviceAddr); tensor.deviceAddr nullptr; } if (tensor.hostAddr ! nullptr) { aclrtFreeHost(tensor.hostAddr); tensor.hostAddr nullptr; } } void DestroyArgs(ArgContext context) { DestroyTensor(context.metadata); DestroyTensor(context.cuSeqLensQOptional); DestroyTensor(context.cuSeqLensKOptional); DestroyTensor(context.seqUsedQOptional); DestroyTensor(context.seqUsedKOptional); DestroyTensor(context.cmpResidualKOptional); if (context.layoutQ ! nullptr) { free(context.layoutQ); context.layoutQ nullptr; } if (context.layoutK ! nullptr) { free(context.layoutK); context.layoutK nullptr; } } aclnnStatus CreateArgs(const ArgScenario scenario, ArgContext context) { ScopeGuard argsGuard([] { DestroyArgs(context); }); aclnnStatus ret; int64_t batchSize 1; context.maxSeqLenQ 16; context.maxSeqLenK 4; context.layoutQ (char *)malloc(sizeof(char) * 16); context.layoutK (char *)malloc(sizeof(char) * 16); strcpy(context.layoutQ, scenario.hasCuSeq ? TND : BSND); strcpy(context.layoutK, scenario.hasCuSeq ? TND : BSND); ret CreateTensor(aclDataType::ACL_INT32, { DLI_METADATA_SIZE }, context.metadata); CHECK_LOG_RET(ret ACL_SUCCESS, ret, Create metadata failed. Error: %d, ret); if (scenario.hasCuSeq) { ret CreateTensor(aclDataType::ACL_INT32, { batchSize 1 }, context.cuSeqLensQOptional); CHECK_LOG_RET(ret ACL_SUCCESS, ret, Create cuSeqLensQOptional failed. Error: %d, ret); ret CreateTensor(aclDataType::ACL_INT32, { batchSize 1 }, context.cuSeqLensKOptional); CHECK_LOG_RET(ret ACL_SUCCESS, ret, Create cuSeqLensKOptional failed. Error: %d, ret); SetInt32TensorData(context.cuSeqLensQOptional, { 0, static_castint32_t(context.maxSeqLenQ) }); SetInt32TensorData(context.cuSeqLensKOptional, { 0, static_castint32_t(context.maxSeqLenK) }); context.batchSize 0; } else { context.batchSize batchSize; } argsGuard.Dismiss(); return ACL_SUCCESS; } void PrintMetadata(const DliGradKLLossMetaData metadata) { printf(coreNum : %d\n, metadata.coreNum); printf(totalSize : %d\n, metadata.totalSize); printf(splitFactorSize : %d\n, metadata.splitFactorSize); for (uint32_t i 0; i std::minuint32_t(metadata.coreNum, DLI_METADATA_MAX_CORE_NUM); i) { printf(bS1Index[%u] : %d\n, i, metadata.bS1Index[i]); } } int main() { int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_LOG_RET(ret ACL_SUCCESS, ret, Init acl failed. ERROR: %d, ret); ScopeGuard sysGuard([] { Finalize(deviceId, stream); }); ArgScenario scenario {}; scenario.hasCuSeq true; ArgContext context {}; ret CreateArgs(scenario, context); CHECK_LOG_RET(ret ACL_SUCCESS, ret, Create input arguments failed. ERROR: %d, ret); ScopeGuard argsGuard([] { DestroyArgs(context); }); uint64_t workspaceSize 0; aclOpExecutor *executor nullptr; void *workspaceAddr nullptr; ret aclnnDenseLightningIndexerKLLossGradMetadataGetWorkspaceSize( context.cuSeqLensQOptional.data, context.cuSeqLensKOptional.data, context.seqUsedQOptional.data, context.seqUsedKOptional.data, context.cmpResidualKOptional.data, context.batchSize, context.maxSeqLenQ, context.maxSeqLenK, context.numHeadsQ, context.numHeadsK, context.headDim, context.layoutQ, context.layoutK, context.maskMode, context.cmpRatio, context.metadata.data, workspaceSize, executor); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclnnDenseLightningIndexerKLLossGradMetadataGetWorkspaceSize failed. ERROR: %d, ret); if (workspaceSize static_castuint64_t(0)) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_LOG_RET(ret ACL_SUCCESS, ret, allocate workspace failed. ERROR: %d, ret); } ScopeGuard workspaceGuard([] { if (workspaceAddr ! nullptr) { aclrtFree(workspaceAddr); workspaceAddr nullptr; } }); ret aclnnDenseLightningIndexerKLLossGradMetadata(workspaceAddr, workspaceSize, executor, stream); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclnnDenseLightningIndexerKLLossGradMetadata failed. ERROR: %d, ret); ret aclrtSynchronizeStream(stream); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclrtSynchronizeStream failed. ERROR: %d, ret); DliGradKLLossMetaData result {}; ret aclrtMemcpy(result, sizeof(result), context.metadata.deviceAddr, sizeof(result), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclrtMemcpy failed. ERROR: %d, ret); PrintMetadata(result); return 0; }示例要点说明示例构造了 TND 场景hasCuSeq true此时cuSeqLensQOptional [0, maxSeqLenQ]、cuSeqLensKOptional [0, maxSeqLenK]batchSize填 0 交由算子通过 cu_seqlens 推导若想走 BSND 场景将hasCuSeq置为 false 即可此时需要显式传入batchSize。完整可编译的示例代码与官方文档示例内容一致仓库中还提供了同名的可执行样例 test_aclnn_dense_lightning_indexer_kl_loss_grad_metadata.cpp 以及 PyTorch 调用样例 test_torch_dense_lightning_indexer_kl_loss_grad_metadata.py可直接参考其完整的Init/CreateArgs/两段式调用/结果回读流程。PyTorch 侧调用方式除 aclnn API 外该算子也可通过 PyTorch 扩展接口调用需先安装 torch_extension 并import cann_ops_transformer。仓库中的 PyTorch 调用样例 展示了 TND GQA KV 压缩的典型用法import torch import cann_ops_transformer metadata torch.ops.cann_ops_transformer.dense_lightning_indexer_kl_loss_grad_metadata( cu_seqlens_qtorch.tensor([0, 123, 230, 234, 511], dtypetorch.int32).npu(), cu_seqlens_ktorch.tensor( [0, 3048, 4098 * 2, 4364 * 3, 4098 * 4], dtypetorch.int32 ).npu(), seqused_qNone, seqused_kNone, cmp_residual_kNone, batch_size4, max_seqlen_q180, max_seqlen_k5, num_heads_q64, num_heads_k1, head_dim128, mask_mode0, layout_qTND, layout_kTND, cmp_ratio13, )该样例展示了 batchSize4、numHeadsQ64GQAgroup64、cmpRatio13KV 压缩的 TND 场景配置与 aclnn 接口在语义上完全对应。问题定位说明关于 AI CPU 算子 Kernel 常见执行问题或异常错误问题定位方法请参考《故障处理》中故障案例集算子执行问题AI CPU 算子 Kernel 执行报错章节文档原文附有外部链接此处不再展开。基于源码可补充的常见排查思路若第一段接口返回ACLNN_ERR_INNER_NULLPTR561103优先检查workspaceSize与executor指针是否有效对应 宿主实现 的空指针检查。若返回ACLNN_ERR_PARAM_INVALID161002按三层校验逐一排查单参数取值numHeadsK/headDim/layout 等、存在性TND 是否传 cu_seqlens、maskMode3 且 cmpRatio≠1 是否传 cmp_residual_k、一致性各 Tensor 维度、数据类型、batch 推导是否一致校验逻辑见 dense_lightning_indexer_kl_loss_grad_metadata_check.h。AICPU 内核阶段还会对 cu_seqlens 的非负性与单调性、seqused/cmp_residual 的非负性做二次校验相关错误日志可参考 arch35 内核实现。总结DenseLightningIndexerKLLossGradMetadata 是一个定位非常明确的前置 metadata 生成算子它在 Ascend 950 系列上把 DenseLightningIndexerKLLossGrad 的 shape/layout/mask/压缩信息转化为固定 64 长度的 INT32 负载切分信息交给主算子作为metadataOptional直接消费。理解它的关键是抓住三点两段式 aclnn 接口的调用顺序、BSND/TND 两种 shape 推导路径的差异以及 metadata 前 6 个字段与主算子侧共享头文件 dense_lightning_indexer_kl_loss_grad_metadata_arch35.h 中常量定义的一一对应关系。接入时建议始终与主算子配合使用并按上文给出的参数取值范围与错误码排查表进行校验。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 在 CANN ops-transformer 中理解 DenseLightningIndexerKLLossGradMetadata：用于负载均衡的分核 metadata 生成器
在 CANN ops-transformer 中理解 DenseLightningIndexerKLLossGradMetadata用于负载均衡的分核 metadata 生成器【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读DenseLightningIndexerKLLossGradMetadata 是 CANN ops-transformer 中专门为aclnnDenseLightningIndexerKLLossGrad算子准备的前置 AICPU 算子。它接收主算子的 shape、layout、mask 与压缩比例信息基于每个 query 的有效序列长度估算负载将 B/S1 合轴后的任务均衡切分到可用 AIC core 上并输出一份固定布局的 INT32 metadata供主算子在 tiling 阶段直接消费从而减少对 host array 的访问。本文围绕该算子的官方文档展开结合仓库内的源码实现、示例与配套算子系统讲解它的定位、接口、参数约束、metadata 布局及调用方式帮助你理解如何在自己的模型中正确接入这条metadata 前置生成的流水线。产品支持情况仅在 Ascend 950 系列上可用该算子在当前仓库中明确标注为仅支持 Ascend 950PR/Ascend 950DT 系列产品文档表格与 README 保持一致其余产品线均不支持产品是否支持Ascend 950PR/Ascend 950DT支持Atlas A3 训练系列产品/Atlas A3 推理系列产品不支持Atlas A2 训练系列产品/Atlas A2 推理系列产品不支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持这与源码中的架构分派逻辑一致。在 dense_lightning_indexer_kl_loss_grad_metadata_aicpu.cpp 中内核首先从上下文中读取soc_version属性再分发到 arch35A5/Ascend 950 系实现而 Host 侧的参数校验 dense_lightning_indexer_kl_loss_grad_metadata_check.h 则先判断socVersion是否包含Ascend950仅在此前提下才执行完整的 A5 参数校验否则只做最基础的空指针检查并直接返回——从源码结构看该算子的完整能力仅面向 Ascend 950 系列。功能定位主算子的前置分核切分器该算子是一个 AICPU 算子接口功能定义如下根据aclnnDenseLightningIndexerKLLossGrad算子的输入 shape、layout、mask 和压缩比例信息计算并输出分核切分 metadata。输出结果可作为aclnnDenseLightningIndexerKLLossGrad算子的metadataOptional输入减少主算子 tiling 阶段对 host array 的访问。该算子不建议单独使用建议与aclnnDenseLightningIndexerKLLossGrad算子配合使用形成完整工作流接收主算子的 shape 信息包括batchSize、maxSeqLenQ、maxSeqLenK、numHeadsQ、numHeadsK、headDim、layout和mask信息。根据每个 query 对应的有效序列长度估算负载并将 B/S1 合轴后的任务均衡切分到可用 AIC core 上。输出 metadata 后后续作为aclnnDenseLightningIndexerKLLossGrad算子的metadataOptional输入使用。从 AICPU 内核实现 dense_lightning_indexer_kl_loss_grad_metadata_aicpu_arch35.h 可以看到核心计算逻辑CalcSplitInfo()遍历每个 batch累加GetS1SeqSize(bIdx)每个 batch 中 query 的有效序列长度得到合轴后的任务总行数totalNum。GenMetadata()按formerCoreProcessNum CeilDiv(totalNum, aicCoreNum)计算上对齐分核粒度remainCoreProcessNum formerCoreProcessNum - 1为下对齐分核粒度remainCoreNum为处理残余粒度的 core 数量usedCoreNum为实际使用的 AIC core 数量最终写入 metadata 的前 5 个字段。序列长度的获取优先级为seqused_q/seqused_k提供的实际参与运算长度 TND 场景下cu_seqlens的差分 标量max_seqlen见 GetS1SeqSize/GetS2SeqSize。内核注册信息位于 dense_lightning_indexer_kl_loss_grad_metadata_aicpu.json该算子运行在DNN_VM_AICPU引擎kernelSo为libtransformer_aicpu_kernels.socomputeCost与workspaceSize均配置为 100属于用户自定义userDefined: True的 AICPU kernel。函数原型与两段式调用每个算子分为两段式接口必须先调用aclnnDenseLightningIndexerKLLossGradMetadataGetWorkspaceSize获取 workspace 大小再调用aclnnDenseLightningIndexerKLLossGradMetadata执行计算。第一段接口原型aclnnStatus aclnnDenseLightningIndexerKLLossGradMetadataGetWorkspaceSize( const aclTensor *cuSeqLensQOptional, const aclTensor *cuSeqLensKOptional, const aclTensor *seqUsedQOptional, const aclTensor *seqUsedKOptional, const aclTensor *cmpResidualKOptional, int64_t batchSize, int64_t maxSeqLenQ, int64_t maxSeqLenK, int64_t numHeadsQ, int64_t numHeadsK, int64_t headDim, char *layoutQ, char *layoutK, int64_t maskMode, int64_t cmpRatio, const aclTensor *metadata, uint64_t *workspaceSize, aclOpExecutor **executor)第二段接口原型aclnnStatus aclnnDenseLightningIndexerKLLossGradMetadata( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)公开 API 的声明位于 aclnn_dense_lightning_indexer_kl_loss_grad_metadata.h宿主侧实现位于 aclnn_dense_lightning_indexer_kl_loss_grad_metadata.cpp第一段接口先对workspaceSize/executor做空指针检查再从当前平台信息中读取aicCoreNum、aivCoreNum与socVersion执行参数校验对 5 个可选 Tensor 入参分别做 Contiguous 处理最后构造执行器并返回 workspace 大小第二段接口通过CommonOpExecutorRun启动异步执行。第一段接口参数说明输入参数前 15 个参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 TensorcuSeqLensQOptionalconst aclTensor*输入表示不同 batch 中 query 的累积 sequence length。支持空 TensorTND 场景下必传并可通过该入参 shape 推导 batch第一个值固定为 0shape 固定为 (B1, )。INT32ND1 维(B1,)×cuSeqLensKOptionalconst aclTensor*输入表示不同 batch 中 key 的累积 sequence length。支持空 TensorTND 场景下必传第一个值固定为 0shape 固定为 (B1, )。INT32ND1 维(B1,)×seqUsedQOptionalconst aclTensor*输入表示不同 batch 中 query 实际参与运算的 sequence length。支持空 Tensorshape 固定为 (B, )。INT32ND1 维(B,)×seqUsedKOptionalconst aclTensor*输入表示不同 batch 中 key 实际参与运算的 sequence length。支持空 Tensorshape 固定为 (B, )。INT32ND1 维(B,)×cmpResidualKOptionalconst aclTensor*输入表示不同 batch 中 key 的 sequence length 与 cmpRatio 相关的残差。支持空 Tensorshape 固定为 (B, )。INT32ND1 维(B,)×batchSizeint64_t输入表示 batch 数量。支持非负数TND 场景可填 0并通过 cuSeqLensQOptional 推导建议值为 0。----maxSeqLenQint64_t输入表示 query 的最大 sequence length。支持非负数BSND 场景必须为正数建议值为 0。----maxSeqLenKint64_t输入表示 key 的最大 sequence length。支持非负数BSND 场景必须为正数建议值为 0。----numHeadsQint64_t输入表示 query 的 head 个数。必须为正数并且能被 numHeadsK 整除当前支持 [1, 128]。----numHeadsKint64_t输入表示 key 的 head 个数。必须为正数当前仅支持 1。----headDimint64_t输入表示 q/k 的 head dimension。必须为正数当前仅支持 128。----layoutQchar*输入表示 query 侧的排列格式。支持 BSND、TND建议值为 BSND。----layoutKchar*输入表示 key 侧的排列格式。支持 BSND、TND建议值为 BSND。----maskModeint64_t输入表示 mask 模式。0: No mask3: rightDownCausal 模式的 mask对应以右顶点为划分的下三角场景建议值为 0。----cmpRatioint64_t输入表示 key 的压缩率。取值范围 [1, 128]建议值 1表示无压缩。----Host 侧校验常量与文档给出的取值范围一一对应在 dense_lightning_indexer_kl_loss_grad_metadata_check.h 中定义了DLI_NO_MASK_MODE 0、DLI_CAUSAL_MASK_MODE 3、DLI_CMP_RATIO_LOWER_BOUND 1、DLI_CMP_RATIO_UPPER_BOUND 128、DLI_NUM_HEADS_Q_LOWER_BOUND_A5 1、DLI_NUM_HEADS_Q_UPPER_BOUND_A5 128以及DLIKG_METADATA_SIZE 64并在CheckSingleParamDli中对numHeadsK 1、headDim 128、layoutQ/layoutK ∈ {BSND, TND}等约束逐一强制校验。输出参数参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 Tensormetadataconst aclTensor*输出表示负载均衡结果输出。输出结果作为 aclnnDenseLightningIndexerKLLossGrad 的 metadataOptional 输入。INT32ND1 维shape 固定为 (64,)×workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程。-----metadata 的 shape 由 infershape 阶段直接写死为 64见 dense_lightning_indexer_kl_loss_grad_metadata_infershape.cpp其中InferShape将输出第 0 维设为optiling::DLIKG_METADATA_SIZE64InferDtype将输出类型固定为DT_INT32。返回值与错误码aclnnStatus返回状态码具体参见 aclnn 返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_INNER_NULLPTR561103workspaceSize 或 executor 为空指针。ACLNN_ERR_PARAM_INVALID多场景161002cuSeqLensQOptional、cuSeqLensKOptional、seqUsedQOptional、seqUsedKOptional 或 cmpResidualKOptional 的数据类型或维度不在支持范围内metadata 为空指针或其数据类型、维度或 shape 不在支持范围内numHeadsQ、numHeadsK 或 headDim 不在支持范围内layoutQ 或 layoutK 为空指针或不在支持范围内batchSize、maxSeqLenQ、maxSeqLenK、maskMode 或 cmpRatio 不满足取值约束TND 场景下未传入对应的 cuSeqLensQOptional 或 cuSeqLensKOptional或 maskMode3 且 cmpRatio 不为 1 时未传入 cmpResidualKOptional从 query、key 或 cmpResidualKOptional 推导出的 batch 数量不一致。从源码看第一段接口的空指针检查在 aclnn_dense_lightning_indexer_kl_loss_grad_metadata.cpp 中实现workspaceSize/executor为空即返回ACLNN_ERR_INNER_NULLPTR而参数范围与一致性校验则由ParamsCheckDli → ParamsCheckDliA5依次完成CheckSingleParamDli单参数取值、CheckExistenceDliTND 必传项与 mask/cmpRatio 联动必传项、CheckConsistencyDli维度、数据类型、batch 一致性三层校验。第二段接口参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnDenseLightningIndexerKLLossGradMetadataGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。返回值同为aclnnStatus具体参见 aclnn 返回码。约束说明确定性说明aclnnDenseLightningIndexerKLLossGradMetadata默认确定性实现。BSND 场景必传batchSize、maxSeqLenQ和maxSeqLenK参数以获取 shape 信息。TND 场景必传cuSeqLensQOptional、cuSeqLensKOptional参数以获取正确 shape 信息当batchSize为 0 时通过cuSeqLensQOptional的 shape 推导 batch。Maskmask_modemask_mode含义备注0无 mask。支持3rightDownCausal 模式的 mask对应以右顶点为划分的下三角场景。支持mask 模式的详细介绍可参考 sparse_mode 参数说明。特殊约束Batch 取值规则如果batchSize大于 0优先使用batchSize。如果batchSize小于等于 0且layoutQ为 TND则通过cuSeqLensQOptional的 shape 推导 batch。如果batchSize小于等于 0且layoutQ为 BSND则报错。Seqlen 取值规则TND 场景下通过cuSeqLensQOptional和cuSeqLensKOptional计算每个 batch 的实际 q/k 长度。BSND 场景下通过maxSeqLenQ和maxSeqLenK获取 q/k 长度。layout 约束layoutQ必须为 BSND 或 TND。layoutK支持 BSND 和 TND建议与layoutQ保持一致。head 约束numHeadsQ、numHeadsK和headDim必须为正数。numHeadsQ必须能被numHeadsK整除。dense 约束cmpRatio取值范围为 [1, 128]。maskMode当前仅支持 0 和 3。上述规则在 AICPU 内核的ParamsInit见 arch35 实现中有对应处理groupSize_ numHeadsQ_ / numHeadsK_计算 group sizemaskMode为 3RIGHT_DOWN_CAUSAL时attentionMode_置 1S2 基准粒度s2BaseSize_在 Ascend 910 系为 2048、Ascend 950 系为 128进一步印证了该算子在 950 系列上的专用适配。Metadata 输出布局metadata 输出为 INT32 Tensor当前 shape 固定为 (64,)字段布局如下字段index说明totalSize0B/S1 合轴后的任务总行数。formerCoreSize1上对齐场景下的分核粒度。remainCoreSize2下对齐场景下的分核粒度。remainCoreNum3处理 remainCoreSize 的 cube core 数量。usedCoreNum4主 kernel 实际使用的 cube core 数量。maxSeqK5key 的最大 sequence length。reserved6-63预留字段当前置 0。该布局与主算子侧共享的元数据定义完全一致在 dense_lightning_indexer_kl_loss_grad_metadata_arch35.h 中optiling命名空间定义了DLIKG_METADATA_SIZE 64、TOTAL_NUM 0、FORMER_CORE_PROCESS_NUM 1、REMAIN_CORE_PROCESS_NUM 2、REMAIN_CORE_NUM 3、USED_CORE_NUM 4等常量AICPU 内核正是按这些索引把负载切分结果写入gradMetadata[]并通过static_assert保证 64 个 INT32 足以承载DlikgMetadata结构。metadata 由前置算子生成、主算子消费前后两个算子通过这份固定协议完成分核切分信息的传递。调用示例以下代码演示了两段式完整调用流程TND 场景示例参数batchSize1、maxSeqLenQ16、maxSeqLenK4、numHeadsQ8、numHeadsK1、headDim128、maskMode0、cmpRatio4具体编译和执行过程请参见编译与运行样例#include algorithm #include cstdint #include cstdio #include cstring #include functional #include iostream #include utility #include vector #include acl/acl.h #include aclnnop/aclnn_dense_lightning_indexer_kl_loss_grad_metadata.h #define CHECK_LOG_RET(cond, ret_val, fmt, ...) \ do { \ if (!(cond)) { \ printf(fmt \n, ##__VA_ARGS__); \ return (ret_val); \ } \ } while (0) constexpr uint32_t DLI_METADATA_MAX_CORE_NUM 25; constexpr uint32_t DLI_METADATA_HEADER_SIZE 8; constexpr uint32_t DLI_METADATA_SIZE 64; struct DliGradKLLossMetaData { int32_t coreNum; int32_t totalSize; int32_t splitFactorSize; int32_t reserved[DLI_METADATA_HEADER_SIZE - 3]; int32_t bS1Index[DLI_METADATA_MAX_CORE_NUM]; }; struct ScopeGuard { explicit ScopeGuard(std::functionvoid() onExitScope) : m_exitFunc(std::move(onExitScope)), m_isDismissed(false) {} ScopeGuard(const ScopeGuard) delete; ScopeGuard operator(const ScopeGuard) delete; ~ScopeGuard() { if (!m_isDismissed) { m_exitFunc(); } } void Dismiss() { m_isDismissed true; } std::functionvoid() m_exitFunc; bool m_isDismissed; }; struct Tensor { void *hostAddr { nullptr }; void *deviceAddr { nullptr }; aclTensor *data { nullptr }; }; struct ArgScenario { bool hasCuSeq { true }; }; struct ArgContext { Tensor cuSeqLensQOptional {}; Tensor cuSeqLensKOptional {}; Tensor seqUsedQOptional {}; Tensor seqUsedKOptional {}; Tensor cmpResidualKOptional {}; Tensor metadata {}; int64_t batchSize { 0 }; int64_t maxSeqLenQ { 0 }; int64_t maxSeqLenK { 0 }; int64_t numHeadsQ { 8 }; int64_t numHeadsK { 1 }; int64_t headDim { 128 }; char *layoutQ { nullptr }; char *layoutK { nullptr }; int64_t maskMode { 0 }; int64_t cmpRatio { 4 }; }; int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } aclnnStatus Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclInit failed. ERROR: %d, ret); ret aclrtSetDevice(deviceId); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclrtSetDevice failed. ERROR: %d, ret); ret aclrtCreateStream(stream); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclrtCreateStream failed. ERROR: %d, ret); return ACL_SUCCESS; } void Finalize(int32_t deviceId, aclrtStream stream) { aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); } aclnnStatus CreateTensor(aclDataType dataType, const std::vectorint64_t shape, Tensor tensor) { auto size GetShapeSize(shape) * aclDataTypeSize(dataType); auto ret aclrtMallocHost((tensor.hostAddr), size); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclrtMallocHost failed. ERROR: %d, ret); memset(tensor.hostAddr, 0, size); ret aclrtMalloc((tensor.deviceAddr), size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclrtMalloc failed. ERROR: %d, ret); tensor.data aclCreateTensor(shape.data(), shape.size(), dataType, nullptr, 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), tensor.deviceAddr); ret aclrtMemcpy(tensor.deviceAddr, size, tensor.hostAddr, size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclrtMemcpy failed. ERROR: %d, ret); return ACL_SUCCESS; } void SetInt32TensorData(Tensor tensor, const std::vectorint32_t hostData) { auto size hostData.size() * sizeof(int32_t); memcpy(tensor.hostAddr, hostData.data(), size); aclrtMemcpy(tensor.deviceAddr, size, tensor.hostAddr, size, ACL_MEMCPY_HOST_TO_DEVICE); } void DestroyTensor(Tensor tensor) { if (tensor.data ! nullptr) { aclDestroyTensor(tensor.data); tensor.data nullptr; } if (tensor.deviceAddr ! nullptr) { aclrtFree(tensor.deviceAddr); tensor.deviceAddr nullptr; } if (tensor.hostAddr ! nullptr) { aclrtFreeHost(tensor.hostAddr); tensor.hostAddr nullptr; } } void DestroyArgs(ArgContext context) { DestroyTensor(context.metadata); DestroyTensor(context.cuSeqLensQOptional); DestroyTensor(context.cuSeqLensKOptional); DestroyTensor(context.seqUsedQOptional); DestroyTensor(context.seqUsedKOptional); DestroyTensor(context.cmpResidualKOptional); if (context.layoutQ ! nullptr) { free(context.layoutQ); context.layoutQ nullptr; } if (context.layoutK ! nullptr) { free(context.layoutK); context.layoutK nullptr; } } aclnnStatus CreateArgs(const ArgScenario scenario, ArgContext context) { ScopeGuard argsGuard([] { DestroyArgs(context); }); aclnnStatus ret; int64_t batchSize 1; context.maxSeqLenQ 16; context.maxSeqLenK 4; context.layoutQ (char *)malloc(sizeof(char) * 16); context.layoutK (char *)malloc(sizeof(char) * 16); strcpy(context.layoutQ, scenario.hasCuSeq ? TND : BSND); strcpy(context.layoutK, scenario.hasCuSeq ? TND : BSND); ret CreateTensor(aclDataType::ACL_INT32, { DLI_METADATA_SIZE }, context.metadata); CHECK_LOG_RET(ret ACL_SUCCESS, ret, Create metadata failed. Error: %d, ret); if (scenario.hasCuSeq) { ret CreateTensor(aclDataType::ACL_INT32, { batchSize 1 }, context.cuSeqLensQOptional); CHECK_LOG_RET(ret ACL_SUCCESS, ret, Create cuSeqLensQOptional failed. Error: %d, ret); ret CreateTensor(aclDataType::ACL_INT32, { batchSize 1 }, context.cuSeqLensKOptional); CHECK_LOG_RET(ret ACL_SUCCESS, ret, Create cuSeqLensKOptional failed. Error: %d, ret); SetInt32TensorData(context.cuSeqLensQOptional, { 0, static_castint32_t(context.maxSeqLenQ) }); SetInt32TensorData(context.cuSeqLensKOptional, { 0, static_castint32_t(context.maxSeqLenK) }); context.batchSize 0; } else { context.batchSize batchSize; } argsGuard.Dismiss(); return ACL_SUCCESS; } void PrintMetadata(const DliGradKLLossMetaData metadata) { printf(coreNum : %d\n, metadata.coreNum); printf(totalSize : %d\n, metadata.totalSize); printf(splitFactorSize : %d\n, metadata.splitFactorSize); for (uint32_t i 0; i std::minuint32_t(metadata.coreNum, DLI_METADATA_MAX_CORE_NUM); i) { printf(bS1Index[%u] : %d\n, i, metadata.bS1Index[i]); } } int main() { int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_LOG_RET(ret ACL_SUCCESS, ret, Init acl failed. ERROR: %d, ret); ScopeGuard sysGuard([] { Finalize(deviceId, stream); }); ArgScenario scenario {}; scenario.hasCuSeq true; ArgContext context {}; ret CreateArgs(scenario, context); CHECK_LOG_RET(ret ACL_SUCCESS, ret, Create input arguments failed. ERROR: %d, ret); ScopeGuard argsGuard([] { DestroyArgs(context); }); uint64_t workspaceSize 0; aclOpExecutor *executor nullptr; void *workspaceAddr nullptr; ret aclnnDenseLightningIndexerKLLossGradMetadataGetWorkspaceSize( context.cuSeqLensQOptional.data, context.cuSeqLensKOptional.data, context.seqUsedQOptional.data, context.seqUsedKOptional.data, context.cmpResidualKOptional.data, context.batchSize, context.maxSeqLenQ, context.maxSeqLenK, context.numHeadsQ, context.numHeadsK, context.headDim, context.layoutQ, context.layoutK, context.maskMode, context.cmpRatio, context.metadata.data, workspaceSize, executor); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclnnDenseLightningIndexerKLLossGradMetadataGetWorkspaceSize failed. ERROR: %d, ret); if (workspaceSize static_castuint64_t(0)) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_LOG_RET(ret ACL_SUCCESS, ret, allocate workspace failed. ERROR: %d, ret); } ScopeGuard workspaceGuard([] { if (workspaceAddr ! nullptr) { aclrtFree(workspaceAddr); workspaceAddr nullptr; } }); ret aclnnDenseLightningIndexerKLLossGradMetadata(workspaceAddr, workspaceSize, executor, stream); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclnnDenseLightningIndexerKLLossGradMetadata failed. ERROR: %d, ret); ret aclrtSynchronizeStream(stream); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclrtSynchronizeStream failed. ERROR: %d, ret); DliGradKLLossMetaData result {}; ret aclrtMemcpy(result, sizeof(result), context.metadata.deviceAddr, sizeof(result), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_LOG_RET(ret ACL_SUCCESS, ret, aclrtMemcpy failed. ERROR: %d, ret); PrintMetadata(result); return 0; }示例要点说明示例构造了 TND 场景hasCuSeq true此时cuSeqLensQOptional [0, maxSeqLenQ]、cuSeqLensKOptional [0, maxSeqLenK]batchSize填 0 交由算子通过 cu_seqlens 推导若想走 BSND 场景将hasCuSeq置为 false 即可此时需要显式传入batchSize。完整可编译的示例代码与官方文档示例内容一致仓库中还提供了同名的可执行样例 test_aclnn_dense_lightning_indexer_kl_loss_grad_metadata.cpp 以及 PyTorch 调用样例 test_torch_dense_lightning_indexer_kl_loss_grad_metadata.py可直接参考其完整的Init/CreateArgs/两段式调用/结果回读流程。PyTorch 侧调用方式除 aclnn API 外该算子也可通过 PyTorch 扩展接口调用需先安装 torch_extension 并import cann_ops_transformer。仓库中的 PyTorch 调用样例 展示了 TND GQA KV 压缩的典型用法import torch import cann_ops_transformer metadata torch.ops.cann_ops_transformer.dense_lightning_indexer_kl_loss_grad_metadata( cu_seqlens_qtorch.tensor([0, 123, 230, 234, 511], dtypetorch.int32).npu(), cu_seqlens_ktorch.tensor( [0, 3048, 4098 * 2, 4364 * 3, 4098 * 4], dtypetorch.int32 ).npu(), seqused_qNone, seqused_kNone, cmp_residual_kNone, batch_size4, max_seqlen_q180, max_seqlen_k5, num_heads_q64, num_heads_k1, head_dim128, mask_mode0, layout_qTND, layout_kTND, cmp_ratio13, )该样例展示了 batchSize4、numHeadsQ64GQAgroup64、cmpRatio13KV 压缩的 TND 场景配置与 aclnn 接口在语义上完全对应。问题定位说明关于 AI CPU 算子 Kernel 常见执行问题或异常错误问题定位方法请参考《故障处理》中故障案例集算子执行问题AI CPU 算子 Kernel 执行报错章节文档原文附有外部链接此处不再展开。基于源码可补充的常见排查思路若第一段接口返回ACLNN_ERR_INNER_NULLPTR561103优先检查workspaceSize与executor指针是否有效对应 宿主实现 的空指针检查。若返回ACLNN_ERR_PARAM_INVALID161002按三层校验逐一排查单参数取值numHeadsK/headDim/layout 等、存在性TND 是否传 cu_seqlens、maskMode3 且 cmpRatio≠1 是否传 cmp_residual_k、一致性各 Tensor 维度、数据类型、batch 推导是否一致校验逻辑见 dense_lightning_indexer_kl_loss_grad_metadata_check.h。AICPU 内核阶段还会对 cu_seqlens 的非负性与单调性、seqused/cmp_residual 的非负性做二次校验相关错误日志可参考 arch35 内核实现。总结DenseLightningIndexerKLLossGradMetadata 是一个定位非常明确的前置 metadata 生成算子它在 Ascend 950 系列上把 DenseLightningIndexerKLLossGrad 的 shape/layout/mask/压缩信息转化为固定 64 长度的 INT32 负载切分信息交给主算子作为metadataOptional直接消费。理解它的关键是抓住三点两段式 aclnn 接口的调用顺序、BSND/TND 两种 shape 推导路径的差异以及 metadata 前 6 个字段与主算子侧共享头文件 dense_lightning_indexer_kl_loss_grad_metadata_arch35.h 中常量定义的一一对应关系。接入时建议始终与主算子配合使用并按上文给出的参数取值范围与错误码排查表进行校验。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考