CANN ops-transformer DenseLightningIndexerKLLossGradMetadata 算子详解：AICPU 分核切分 Metadata 生成原理与调用实践

CANN ops-transformer DenseLightningIndexerKLLossGradMetadata 算子详解AICPU 分核切分 Metadata 生成原理与调用实践【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformerDenseLightningIndexerKLLossGradMetadata 是 CANN ops-transformer 仓库中为DenseLightningIndexerKLLossGradLightningIndexer KL Loss 反向、Dense 模式算子配套的前置 AICPU 算子。它负责根据主算子的输入 shape、layout、mask 与压缩比例cmp_ratio在 Host/CPU 侧预先完成 B/S1 合轴任务的负载均衡切分并输出固定 64 长度的分核 metadata供主算子核侧直接使用。读完本文你将掌握该算子的设计动机、完整参数语义、AICPU 分核切分的数学原理、aclnn 两段式接口与 PyTorch 调用的完整实践方法。算子定位为反向算子预计算的调度前置算子DenseLightningIndexerKLLossGrad 是 LightningIndexer 分支 KL Loss 的反向计算算子Dense 模式接收 q、k、w 以及主 Attention 分支预计算的 attnSoftmaxL1Norm 与 softmaxLse输出 dq、dk、dw 与 softmaxOut。其核心公式包括 Indexer logitsS qK^T、I WReLU(S)、Softmax 输出 y、KL 散度L(I) Σ_t D_KL(p_t,: || Softmax(I_t,:))以及梯度表达式dI Softmax(I) - p等详见 主算子 README。由于主算子在 NPU 上以多核并行方式执行需要在 tiling 阶段确定哪个核处理哪些任务。传统做法是在主算子 tiling 中实时访问 Host 侧 array 完成分核计算而 DenseLightningIndexerKLLossGradMetadata 的职责就是把这部分工作提前到独立的前置算子中完成接收主算子的 shape 信息batchSize、maxSeqLenQ、maxSeqLenK、numHeadsQ、numHeadsK、headDim、layout 和 mask 信息根据每个 query 对应的有效序列长度估算负载将 B/S1 合轴后的任务均衡切分到可用的 AIC 核上输出 metadata作为DenseLightningIndexerKLLossGrad算子的metadataOptional输入减少主算子 tiling 阶段对 host array 的访问。因此该算子不建议单独使用必须与主算子配套形成完整工作流主算子侧也将 metadata 声明为必须传入的输入用于计算分核信息见 主算子约束说明。从仓库结构可以推断metadata 数据结构DlikgMetadata定义在主算子侧的 dense_lightning_indexer_kl_loss_grad_metadata_arch35.h而生成逻辑位于本算子目录两侧共享同一份常量定义DLIKG_METADATA_SIZE 64、GRAD_METADATA_SIZE 6并通过static_assert保证 64 个 int32 足以容纳全部字段。产品支持情况产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品×Atlas A2 训练系列产品 / Atlas A2 推理系列产品×Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×从 AICPU kernel 源码dense_lightning_indexer_kl_loss_grad_metadata_aicpu_arch35.h也可以印证kernel 按soc_version属性分流ValidSocVersion枚举仅区分 Ascend910 与 Ascend950 两个分支而 Host 侧校验dense_lightning_indexer_kl_loss_grad_metadata_check.h会先判断socVersion是否包含 Ascend950仅在 Ascend950 系列上执行完整的 A5 参数校验其他平台直接返回空校验通过与该产品支持表格完全一致。完整参数说明算子共包含 5 个可选输入、9 个属性/可选属性与 1 个输出全部罗列如下对应 README 的参数表。参数名输入/输出/属性描述数据类型数据格式cu_seqlens_q可选输入不同 Batch 中 Query 的累积有效 Sequence Lengthshape 为 (B1,)仅 layout_q 为 TND 场景下必传第一个值固定为 0INT32NDcu_seqlens_k可选输入不同 Batch 中 Key 的累积有效 Sequence Lengthshape 为 (B1,)仅 layout_k 为 TND 场景下必传第一个值固定为 0INT32NDseqused_q可选输入不同 Batch 中 Query 实际参与运算的 Sequence Lengthshape 为 (B,)INT32NDseqused_k可选输入不同 Batch 中 Key 实际参与运算的 Sequence Lengthshape 为 (B,)INT32NDcmp_residual_k可选输入不同 Batch 中 cmp_kv 压缩后 Sequence Length 的余数shape 为 (B,)配合 cmp_ratio 实现 cmp_kv 部分的 mask 与负载计算cmp_ratio 不为 1 且 mask_mode 为 3 时必传INT32NDnum_heads_q属性Query 的 head 个数当前支持 [1, 128]INT32-num_heads_k属性Key 的 head 个数当前仅支持 1INT32-head_dim属性注意力头维度当前仅支持 128INT32-batch_size可选属性Batch 数量默认值为 0INT32-max_seqlen_q可选属性Query 的最长 Sequence Length默认值为 0INT32-max_seqlen_k可选属性Key 的最长 Sequence Length默认值为 0INT32-layout_q可选属性Query 排列格式支持 BSND、TND默认值为 BSNDSTRING-layout_k可选属性Key 排列格式支持 BSND、TND默认值为 BSNDSTRING-mask_mode可选属性mask 模式0 表示 No mask3 表示 rightDownCausal 模式默认值为 0INT32-cmp_ratio可选属性Key 的压缩率取值范围 [1, 128]默认值为 1无压缩INT32-metadata输出负载均衡结果输出shape 固定为 [64]INT32ND关键参数的语义边界cu_seqlens_q / cu_seqlens_k要求元素为当前 Batch 与前序 Batch 有效 token 数的累加值且后一个元素必须大于等于前一个元素递增序列。AICPU kernel 的ParamsCheck会逐元素校验非负性与单调性不满足时返回KERNEL_STATUS_PARAM_INVALID并打印具体越界位置见 源码。seqused_q / seqused_k表示每个 Batch 中实际参与运算的有效 token 数实际序列长度shape 为 (B,)所有元素要求非负。若传入kernel 将优先使用它推导 batch 与各 Batch 的序列长度。cmp_residual_k需满足cmp_residual_k[i] cmp_ratio用于在 Key 侧做了 cmp 压缩cmp_ratio ! 1且叠加 rightDownCausal maskmask_mode 3时还原压缩前的 Key 长度以正确计算 mask 与负载。batch_size / max_seqlen_q / max_seqlen_k默认值为 0且取值遵循建议值为 0、通过 Tensor shape 推导的原则——Batch 优先从seqused_q的 shape 推导其次在 TND 下从cu_seqlens_q的 shape 推导dim0 - 1最后才回退到batch_size属性对应 kernel 中GetQueryBatchSize的实现。mask_mode 语义0为无 maskdefaultMask3为 rightDownCausal 模式对应以右顶点为划分的下三角场景。其余 mask 模式allMask、leftUpCausal、band、prefix、global、dilated、block_local在 sparse_mode 参数说明 中有完整介绍但当前版本主算子的约束表明确标注不支持本 metadata 算子同样仅接受 0 与 3。源码级原理AICPU 分核切分的完整执行链路算子注册与运行框架从 aicpu.json 可以看到该算子注册在DNN_VM_AICPU引擎上kernelSo为libtransformer_aicpu_kernels.so同步执行flagAsyncFalse。入口在 dense_lightning_indexer_kl_loss_grad_metadata_aicpu.cpp通过REGISTER_CPU_KERNEL注册DenseLightningIndexerKLLossGradMetadataCpuKernel读取soc_version属性后委托给arch35Kernel_.Compute(ctx)执行。Compute 的三段式流程arch35 实现dense_lightning_indexer_kl_loss_grad_metadata_aicpu_arch35.h将计算组织为Prepare → BalanceSchedule → GenMetadatainline bool DenseLightningIndexerKLLossGradMetadataCpuKernelArch35::Compute(CpuKernelContext ctx) { bool success Prepare(ctx); if (!success) { return KERNEL_STATUS_PARAM_INVALID; } success BalanceSchedule() GenMetadata(); return success ? KERNEL_STATUS_OK : KERNEL_STATUS_PARAM_INVALID; }Prepare绑定 5 个输入 TensorcuSeqlensQ_、cuSeqlensK_、sequsedQ_、sequsedK_、cmpResidualK_与输出metadata_读取必选属性aic_core_num、aiv_core_num、soc_version、num_heads_q、num_heads_k、head_dim再以可选方式读取batch_size、max_seqlen_q、max_seqlen_k、layout_q、layout_k、mask_mode、cmp_ratio随后执行ParamsCheck合法性校验与ParamsInit初始化派生参数。BalanceSchedule调用CalcSplitInfo遍历所有 Batch累加每个 Batch 的 Query 侧有效序列长度得到 B/S1 合轴后的任务总行数totalNum对应 metadata 字段totalSize。GenMetadata基于totalNum与可用 AIC 核数完成负载均衡切分并写入 metadata 输出。其中ParamsInit还完成了几项关键的派生计算groupSize_ num_heads_q / num_heads_k即 GQA/MQA 场景下每组 KV 对应的 Query head 数当前 num_heads_k 固定为 1因此 groupSize 恒等于 num_heads_q根据soc_version设置s2BaseSize_Ascend910 为 2048Ascend950 为 128核内 S2 维度的基础切分粒度当前代码中作为预留的基准尺寸根据mask_mode设置attentionMode_与preToken_rightDownCausal3与 band 模式映射为 attentionMode1causal 类无 mask0映射为 attentionMode0。分核切分算法GenMetadata的核心逻辑如下源码 L324-L339uint32_t formerCoreProcessNum CeilDiv(totalNum, aicCoreNum_); // 上对齐分核粒度 uint32_t remainCoreProcessNum formerCoreProcessNum - 1; // 下对齐分核粒度 uint32_t remainCoreNum formerCoreProcessNum * aicCoreNum_ - totalNum; // 处理余数的核数 uint32_t usedCoreNum totalNum aicCoreNum_ ? totalNum : aicCoreNum_; // 实际使用核数其中CeilDiv(num, rnd) (num rnd - 1) / rnd见 公共头文件且对 rnd 为 0 做了保护。这一组公式本质上是一个标准的大核均摊 余数核兜底的均衡切分策略若任务总行数totalNum能被核数整除则每个核分得formerCoreProcessNum行remainCoreNum 0若不能整除则前aicCoreNum_ - remainCoreNum个核各处理formerCoreProcessNum行上对齐剩余remainCoreNum个核各处理formerCoreProcessNum - 1行下对齐从而把负载差控制在 1 行以内实现均衡当任务数少于核数时只启动totalNum个核usedCoreNum totalNum避免空转。Metadata 输出布局metadata 为 INT32 Tensorshape 固定为 (64,)由 infershape 实现 通过IMPL_OP_INFERSHAPE注册输出 shape 与DT_INT32类型。字段布局如下字段index说明totalSize0B/S1 合轴后的任务总行数 totalNumformerCoreSize1上对齐场景下的分核粒度 formerCoreProcessNumremainCoreSize2下对齐场景下的分核粒度 remainCoreProcessNumremainCoreNum3处理 remainCoreSize 的 cube core 数量 remainCoreNumusedCoreNum4主 kernel 实际使用的 cube core 数量 usedCoreNummaxSeqK5Key 的最大 sequence lengthreserved6-63预留字段当前置 0字段常量TOTAL_NUM、FORMER_CORE_PROCESS_NUM、REMAIN_CORE_PROCESS_NUM、REMAIN_CORE_NUM、USED_CORE_NUM在 dense_lightning_indexer_kl_loss_grad_metadata_arch35.h 中定义AICPU kernel 写入时使用的正是同一组索引主算子核侧则通过GetAttrAbsIndex(coreIdx, metaIdx) GRAD_METADATA_SIZE * coreIdx metaIdx按核索引读取分片信息。64 个 int32 中除前 6 个字段外的 58 个预留位置为后续扩展核数或分核策略预留了空间示例代码中DliGradKLLossMetaData结构体以coreNum/totalSize/splitFactorSize reserved[5] bS1Index[25]的形式解读该内存也是一种兼容视角。aclnn 两段式接口调用aclnn 接口遵循 CANN 的两段式接口规范先调用 GetWorkspaceSize 版本获取 workspace 大小并完成入参校验再调用执行版本真正下发计算。函数原型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);第一段接口参数要点参数关键约束cuSeqLensQOptional / cuSeqLensKOptional支持空 TensorTND 场景必传并可通过其 shapedim0 - 1推导 batch第一个值固定为 0shape 固定 (B1,)INT32、ND、不支持非连续 TensorseqUsedQOptional / seqUsedKOptional支持空 Tensorshape 固定 (B,)INT32cmpResidualKOptional支持空 Tensorshape 固定 (B,)INT32batchSize支持非负数TND 场景可填 0 并借助 cuSeqLensQOptional 推导建议值 0maxSeqLenQ / maxSeqLenK支持非负数BSND 场景必须为正数建议值 0numHeadsQ必须为正数且能被 numHeadsK 整除支持 [1, 128]numHeadsK必须为正数当前仅支持 1headDim必须为正数当前仅支持 128layoutQ / layoutK支持 BSND、TND建议 BSNDmaskMode0No mask或 3rightDownCausal以右顶点为划分的下三角场景cmpRatio取值范围 [1, 128]建议值 1无压缩metadata输出INT321 维 shape 固定 (64,)结果作为主算子 metadataOptional 输入workspaceSize / executor输出分别返回 Device 侧 workspace 大小与包含算子计算流程的 op 执行器返回值与常见错误码两段接口返回aclnnStatus完整状态码说明见 aclnn 返回码文档。第一段接口入参校验失败时主要报错场景如下返回值错误码描述ACLNN_ERR_INNER_NULLPTR561103workspaceSize 或 executor 为空指针ACLNN_ERR_PARAM_INVALID161002① 各可选输入的数据类型或维度不在支持范围② metadata 为空指针或其类型/维度/shape 不支持③ numHeadsQ、numHeadsK、headDim 越界④ layoutQ/layoutK 为空指针或不支持⑤ batchSize、maxSeqLenQ、maxSeqLenK、maskMode、cmpRatio 不满足取值约束⑥ TND 下未传对应 cuSeqLens或 maskMode3 且 cmpRatio≠1 时未传 cmpResidualKOptional⑦ 从 query、key、cmpResidualKOptional 推导出的 batch 数量不一致上述校验逻辑在 Host 侧 dense_lightning_indexer_kl_loss_grad_metadata_check.h 中分为三个层次实现CheckSingleParamDli单参数范围、CheckExistenceDliTND 必传项与 cmp 场景必传项、CheckConsistencyDliTensor 维度/类型/元素个数/跨来源 batch 一致性并在ParamsCheckDli中按 Ascend950 平台门控。完整调用示例仓库提供了可编译运行的样例 test_aclnn_dense_lightning_indexer_kl_loss_grad_metadata.cpp编译与运行方法参见编译与运行样例。以下为核心流程节选#include acl/acl.h #include aclnnop/aclnn_dense_lightning_indexer_kl_loss_grad_metadata.h 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]; };主流程分五步初始化 ACL 环境aclInit→aclrtSetDevice→aclrtCreateStream→ 构造输入/输出 Tensor含CreateTensor完成 Host/Device 内存分配与 H2D 拷贝SetInt32TensorData填充 cu_seqlens 等数据→ 调用第一段接口获取 workspace 并分配 → 调用第二段接口执行 →aclrtSynchronizeStream同步后通过 D2H 拷贝回读 metadata 并打印。以样例中的 TND 场景为例batchSize1、maxSeqLenQ16、maxSeqLenK4、numHeadsQ8、numHeadsK1、headDim128、maskMode0、cmpRatio4cuSeqLensQ [0, 16]、cuSeqLensK [0, 4]batchSize显式置 0 以验证由 cu_seqlens 推导 batch的路径。样例同时通过ArgScenario.hasCuSeq开关覆盖 BSND 场景此时必传batchSize/maxSeqLenQ/maxSeqLenK。PyTorch 调用方式该算子也通过 torch_extension 暴露了 PyTorch 接口样例见 test_torch_dense_lightning_indexer_kl_loss_grad_metadata.pyimport 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, )该示例是一个 batch4 的 TND 变长场景cu_seqlens_q 展示出每个 Batch 的 Query 长度分别为 123、107、4、277累加值 0→123→230→234→511key 侧每个 Batch 长度分别为 3048、4098、4364、4098均大于 Query符合反向算子中每个 query 需要与全部 key 做 dense 注意力的语义max_seqlen_q180等属性在此场景仅作为辅助信息实际序列长度以 cu_seqlens 为准cmp_ratio13表示 key 侧按 13:1 压缩由于 mask_mode0 无 mask无需传 cmp_residual_k。更多 PyTorch 接口细节见 torchapi 文档。使用约束与确定性必须与DenseLightningIndexerKLLossGrad算子配套使用输出直接作为主算子的 metadataOptional 输入BSND 场景必传 batchSize、maxSeqLenQ、maxSeqLenK 获取 shape 信息TND 场景必传 cuSeqLensQOptional、cuSeqLensKOptional当 batchSize 为 0 时通过 cuSeqLensQOptional 的 shape 推导 batch若 batchSize ≤ 0 且 layoutQ 为 BSND则直接报错Batch 取值规则batchSize 0 时优先使用 batchSizebatchSize ≤ 0 且 TND 时由 cu_seqlens_q 推导Seqlen 取值规则TND 下通过 cu_seqLens 差分计算各 Batch 实际 q/k 长度BSND 下直接取 maxSeqLenQ / maxSeqLenKlayout 约束layoutQ 必须为 BSND 或 TNDlayoutK 支持 BSND 与 TND建议与 layoutQ 保持一致head 约束numHeadsQ、numHeadsK、headDim 必须为正数且 numHeadsQ 必须能被 numHeadsK 整除dense 约束cmpRatio ∈ [1, 128]maskMode 当前仅支持 0 和 3确定性该算子为默认确定性实现重复执行同一输入会得到一致的切分结果所有输入 Tensor 均要求非连续 Tensor 不支持×shape、数据类型必须与上表严格一致。问题定位说明关于 AI CPU 算子 Kernel 常见执行问题或异常错误的定位方法可参考 CANN《故障处理》中故障案例集 算子执行问题 AI CPU 算子 Kernel 执行报错一节。此外由于本算子计算逻辑全部在 AICPU 侧完成若出现参数校验类错误kernel 会通过KERNEL_LOG_ERROR输出具体越界元素的下标与数值如cu_seqlens_q[i]为负、非递增序列等可结合 Host 侧ParamsCheckDli的OP_LOGE日志共同定位错误信息中通常会同时包含字段名与期望的取值范围。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考