CANN ops-transformer CompressorGrad 反向算子技术指南aclnnCompressorGrad 接口、梯度公式与 Ascend 950 实现解析【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer本文是 CANN ops-transformer 算子库中CompressorGrad算子aclnnCompressorGrad的完整技术指南。CompressorGrad是 KV 压缩算子Compressor的反向实现用于在压缩注意力Compressed Attention场景下反向计算输入 $X$、KV 压缩权重 $W^{KV}$/$W^{Gate}$ 与位置编码 $Ape$ 的梯度。通过本文读者可以掌握该算子的数学原理、两段式 aclnn 调用流程、全部入参出参的 shape/dtype 约束、典型调用示例以及算子定义、shape 推导、tiling 校验与 kernel 三阶段流水线等仓库源码级实现细节。产品支持情况CompressorGrad算子当前仅支持 Ascend 950 系列产品具体支持矩阵如下与 算子 README 及 接口文档 一致产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持在算子注册层面这一点也有对应体现attention/compressor_grad/op_host/compressor_grad_def.cpp中通过this-AICore().AddConfig(ascend950, aicore_config)仅向 ascend950 平台注册了 AICore 配置tiling 实现中默认SocVersion即为ASCEND950参见attention/compressor_grad/op_host/compressor_grad_tiling.h。功能说明算子定位与输入来源CompressorGrad是Compressor算子的反向算子用于计算输入 $X$、权重 $W^{KV}$/$W^{Gate}$ 与位置编码 $Ape$ 的梯度。前向算子Compressor在gradEnabled为 true 时会额外导出两组中间结果作为本算子的输入softmax_score分组 softmax 结果公式中的 $S$FLOAT32 类型kvsoftmax 结果与kv_state的 Hadamard 乘积中间结果公式中的 $K$FLOAT32 类型。反向输入dCmpKv公式中的 $dC$则是前向输出cmp_kv的上游梯度。主要计算过程算子的反向计算分为四步逐块计算 Hadamard 积反向将上游梯度 $dC$ 与softmax_score、kv逐元素相乘得到 $dK$ 与 $dS^\prime$softmax 反向对 $dS^\prime$ 沿压缩轴做 softmax 反向得到 $dZ$APE 梯度计算按 token 位置累加 $dZ$得到 $dApe$matmul 反向将 $dK$、$dZ$ 与权重做矩阵乘法反向得到 $dX$、$dW^{KV}$、$dW^{Gate}$。计算公式设 $N$ 为压缩块总数$i$ 为压缩块序号$dC_i$ 为第 $i$ 块的上游梯度$S_i$、$K_i$ 分别为softmax_score、kv第 $i$ 块① Hadamard 乘积反向$$ dK_i dC_i \odot S_i,~ i1,\cdots,N $$$$ dS^\prime_i dC_i \odot K_i,~ i1,\cdots,N $$② softmax 反向沿压缩轴求和$k$ 为块内行序号$$ dZ_i S_i \odot \left(dS^\prime_i - \sum_{k1}^{coff \cdot cmp_ratio} \left(S_i \odot dS^\prime_i\right)_{k,:}\right),~ i1,\cdots,N $$③ APE 梯度$pos$ 为 $dZ$ 各行对应 token 的全局位置$$ dApe ScatterAdd\left(dZ,~ pos % cmp_ratio\right) $$④ 矩阵乘法反向$dNewKv$、$dNewScore$ 为 $dK$、$dZ$ 按压缩块映射回全局 token 行的结果coff2时 prev/cur 半区分别对应上一块与本块的 token 行与正向的 $W^{aKV}$/$W^{bKV}$ 对应$$ dX dNewKv W^{KV} dNewScore W^{Gate} $$$$ dW^{KV} dNewKv^T X,~ dW^{Gate} dNewScore^T X $$函数原型每个算子分为两段式接口必须先调用aclnnCompressorGradGetWorkspaceSize接口获取入参并根据流程计算所需 workspace 大小再调用aclnnCompressorGrad接口执行计算。aclnnStatus aclnnCompressorGradGetWorkspaceSize( const aclTensor *x, const aclTensor *wkv, const aclTensor *wgate, const aclTensor *dCmpKv, const aclTensor *softmaxScore, const aclTensor *kv, const aclTensor *cuSeqlensOptional, const aclTensor *sequsedOptional, const aclTensor *startPosOptional, int64_t cmpRatio, int64_t coff, const aclTensor *dXOut, const aclTensor *dWkvOut, const aclTensor *dWgateOut, const aclTensor *dApeOut, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnCompressorGrad( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)aclnnCompressorGradGetWorkspaceSize 参数说明第一段接口完成入参校验并根据算子流程计算所需 workspace 大小。完整参数说明如下参数名输入/输出描述使用说明数据类型数据格式维度shape非连续Tensorx输入公式中的 $X$前向输入的原始数据。不支持空Tensor。FLOAT16、BFLOAT16NDBS合轴[T,H]、BS非合轴[B,S,H]×wkv输入公式中的 $W^{KV}$前向kv压缩权重。不支持空Tensor。FLOAT16、BFLOAT16ND[coff*D,H]×wgate输入公式中的 $W^{Gate}$前向gate压缩权重。不支持空Tensor。FLOAT16、BFLOAT16ND[coff*D,H]×dCmpKv输入公式中的 $dC$前向输出cmp_kv的上游梯度。不支持空Tensor。FLOAT16、BFLOAT16NDBS合轴[min(T,T//cmp_ratioB),D]、BS非合轴[B,ceil(S/cmp_ratio),D]×softmaxScore输入公式中的 $S$前向在gradEnabled为true时导出的分组softmax中间结果。不支持空Tensor。FLOAT32NDBS合轴[min(T,T//cmp_ratioB), coff*cmp_ratio, D]、BS非合轴[B,ceil(S/cmp_ratio),coff*cmp_ratio,D]×kv输入公式中的 $K$前向在gradEnabled为true时导出的softmax结果与kv_state的Hadamard乘积中间结果。不支持空Tensor。FLOAT32ND同softmaxScore×cuSeqlensOptional可选输入表示不同Batch中的有效token数。当x的shape为[T,H]时必传输入shape为[B1,]当x的shape为[B,S,H]时参数必须为空。不支持空Tensor。INT32ND当x的shape为[T,H]时输入shape为[B1,]×sequsedOptional可选输入表示不同Batch中实际参与压缩的token数。为None时表示和每个Batch上的Sequence Length长度相同要求seqused[n]不超过对应Sequence Length且不小于0。INT32ND[B,]×startPosOptional可选输入表示计算起始位置。为None时表示从0开始进行计算。INT32ND[B,]×cmpRatio输入用于稀疏计算表示数据压缩率与前向一致。取值范围为[2, 128]内的整数。INT32---coff可选输入表示是否进行overlap数据重排与前向一致。取值范围为[1, 2]。当coff1时无需进行overlap数据重排。当coff2时需要进行overlap数据重排。INT32---dXOut输出公式中的 $dX$输入x的梯度。不支持空Tensor。FLOAT16、BFLOAT16ND与x相同BS合轴[T,H]、BS非合轴[B,S,H]×dWkvOut输出公式中的 $dW^{KV}$权重wkv的梯度。不支持空Tensor。FLOAT16、BFLOAT16ND[coff*D,H]×dWgateOut输出公式中的 $dW^{Gate}$权重wgate的梯度。不支持空Tensor。FLOAT16、BFLOAT16ND[coff*D,H]×dApeOut输出公式中的 $dApe$APE位置编码的梯度。不支持空Tensor。FLOAT32ND[cmp_ratio,coff*D]×返回值aclnnStatus返回状态码具体参见 aclnn返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001必须传入的参数如接口核心依赖的输入/输出参数中存在空指针。ACLNN_ERR_PARAM_INVALID161002输入参数的shape维度/尺寸、dtype数据类型不在接口支持的范围内。ACLNN_ERR_RUNTIME_ERROR361001API内存调用NPU Runtime接口时发生异常如Runtime服务未启动、内存申请失败等。ACLNN_ERR_INNER_TILING_ERROR561002tiling发生异常入参的dtype类型或者shape错误。关键说明为什么要求所有 Tensor 非空与正向Compressor算子不同CompressorGrad所有输入/输出均不支持空 TensorshapeSize 必须大于 0。这一点在 tiling 实现中有专门校验attention/compressor_grad/op_host/compressor_grad_tiling.cpp的CheckEmptyTensor会逐一检查 x、wkv、wgate、dCmpKv、softmaxScore、kv 及四个输出的GetShapeSize()任一为 0 即直接拦截并报错。注释中明确说明与正向不同正向 x 支持 B/S/T0 走 EMPTY_X 分支反向无空 tensor 分支。aclnnCompressorGrad 参数说明第二段接口基于第一段接口计算得到的 workspace 与 executor 执行实际计算参数名输入/输出描述workspace输入在Device侧申请的workspace内存地址。workspaceSize输入在Device侧申请的workspace大小由第一段接口aclnnCompressorGradGetWorkspaceSize获取。executor输入op执行器包含了算子计算流程。stream输入指定执行任务的Stream。返回值aclnnStatus返回状态码具体参见 aclnn返回码。约束说明维度含义x 参数维度含义BBatch Size表示输入样本批量大小、SSequence Length表示输入样本序列长度、HHead Size表示 hidden 层的大小、DHead Dim表示 hidden 层的最小单元大小、T 表示所有 Batch 输入样本序列长度的累加和。输入 shape 限制wkv 支持输入 shape[coff*D,H]wgate 支持输入 shape[coff*D,H]softmaxScore 支持输入 shapeBS合轴时为[min(T,T//cmp_ratioB), coff*cmp_ratio, D]BS非合轴时为[B,ceil(S/cmp_ratio),coff*cmp_ratio,D]kv 支持输入 shape同 softmaxScoredCmpKv 支持输入 shapeBS合轴时为[min(T,T//cmp_ratioB),D]BS非合轴时为[B,ceil(S/cmp_ratio),D]dX 支持输出 shape与 x 相同BS合轴时为[T,H]、BS非合轴时为[B,S,H]dWkv、dWgate 支持输出 shape[coff*D,H]dApe 支持输出 shape[cmp_ratio,coff*D]startPos 支持输入 shape[B,]。BS 合轴场景x 的 shape 为 [T,H]cuSeqlens 输入 shape 必须为[B1,]。该参数中每个元素的值表示当前 batch 与之前所有 batch 的 token 数总和即前缀和因此后一个元素的值必须大于等于前一个元素的值且第一位必须为 0seqused 支持输入 shape[B,]要求每个 Batch 的有效 token 数小于等于对应 Sequence Length即seqused[n] cu_seqlens[n1] - cu_seqlens[n]且不小于 0。BS 非合轴场景x 的 shape 为 [B,S,H]cuSeqlens 参数必须为空seqused 支持输入 shape[B,]要求每个 Batch 的有效 token 数小于等于对应 Sequence Length即seqused[n] S且不小于 0。输入值域限制该接口支持 B、S 泛化并存在如下场景限制不支持 B、S、T 为 0 的空 Tensor与正向Compressor不同CompressorGrad所有输入/输出均不支持空 TensorshapeSize 必须大于 0部分长序列场景下如果计算量过大可能会导致出现超过 NPU 内存的报错。注意这里计算量受 x 输入 shape 的影响值越大计算量越大。输入属性限制支持 D 为 128/512支持 H 为 1K~10K512 对齐支持 coff 为 1/2支持 cmp_ratio 为 2~128。这些限制与源码中的常量定义一致attention/compressor_grad/op_host/compressor_grad_tiling.h定义了MAX_HIDDEN_SIZE 10240、MIN_HIDDEN_SIZE 1024、ALIGN_FACTOR_HIDDEN_SIZE 512、MIN_CMP_RATIO 2、MAX_CMP_RATIO 128、HEAD_DIM {128, 512}、COFF {1, 2}tiling 的CheckFeature会对 hiddenSize 做 [1024, 10240] 区间与 512 对齐校验CheckSingleParaCmpRatio/CheckSingleParaCoff分别对 cmp_ratio 区间与 coff 取值做校验。确定性计算aclnnCompressorGrad 默认为确定性实现。调用示例示例代码如下具体编译和执行过程请参考 编译与运行样例。完整可编译样例位于 examples/test_aclnn_compressor_grad.cpp。/** * \file test_aclnn_compressor_grad.cpp * \brief CompressorGrad 算子 aclnn 调用示例A5 / ascend950 * 场景C4AD512, coff2, cmp_ratio4, BSH layout, BF16 */ #include algorithm #include cmath #include cstdint #include iostream #include numeric #include vector #include acl/acl.h #include aclnnop/aclnn_compressor_grad.h #include opdev/bfloat16.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) namespace { using op::bfloat16; int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } int Init(int32_t deviceId, aclrtContext *context, aclrtStream *stream) { auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateContext(context, deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateContext failed. ERROR: %d\n, ret); return ret); ret aclrtSetCurrentContext(*context); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetCurrentContext failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void **deviceAddr, aclDataType dataType, aclTensor **tensor) { auto size GetShapeSize(shape) * sizeof(T); if (size 0) { auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); } else { *deviceAddr nullptr; } std::vectorint64_t strides(shape.size(), 1); for (int64_t i static_castint64_t(shape.size()) - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } void PrintBf16Result(const std::vectorint64_t shape, void **deviceAddr) { auto size GetShapeSize(shape); std::vectorbfloat16 resultData(size, bfloat16(0.0f)); auto ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return); for (int64_t i 0; i size i 10; i) { // 10: max print LOG_PRINT(result[%ld] is: %f\n, i, static_castfloat(resultData[i])); } } void PrintF32Result(const std::vectorint64_t shape, void **deviceAddr) { auto size GetShapeSize(shape); std::vectorfloat_t resultData(size, 0.0f); auto ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return); for (int64_t i 0; i size i 10; i) { // 10: max print LOG_PRINT(result[%ld] is: %f\n, i, resultData[i]); } } } // namespace int main() { // 1. device/stream 初始化 int32_t deviceId 0; aclrtContext context nullptr; aclrtStream stream nullptr; auto ret Init(deviceId, context, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 场景参数: C4A (D512, coff2, cmp_ratio4, BSH layout, BF16) int64_t B 1; int64_t S 4; int64_t hiddenSize 4096; int64_t headDim 512; int64_t coff 2; int64_t cmpRatio 4; int64_t coffD coff * headDim; int64_t Sr (S cmpRatio - 1) / cmpRatio; // 2. 构造输入与输出 shape std::vectorint64_t xShape {B, S, hiddenSize}; std::vectorint64_t wkvShape {coffD, hiddenSize}; std::vectorint64_t wgateShape {coffD, hiddenSize}; std::vectorint64_t dCmpKvShape {B, Sr, headDim}; std::vectorint64_t softmaxScoreShape {B, Sr, coff * cmpRatio, headDim}; std::vectorint64_t kvShape {B, Sr, coff * cmpRatio, headDim}; std::vectorint64_t startPosShape {B}; std::vectorint64_t dXShape {B, S, hiddenSize}; std::vectorint64_t dWkvShape {coffD, hiddenSize}; std::vectorint64_t dWgateShape {coffD, hiddenSize}; std::vectorint64_t dApeShape {cmpRatio, coffD}; // 3. 构造 host 数据 int64_t xSize GetShapeSize(xShape); int64_t wkvSize GetShapeSize(wkvShape); int64_t wgateSize GetShapeSize(wgateShape); int64_t dCmpKvSize GetShapeSize(dCmpKvShape); int64_t softmaxScoreSize GetShapeSize(softmaxScoreShape); int64_t kvSize GetShapeSize(kvShape); int64_t dXSize GetShapeSize(dXShape); int64_t dWkvSize GetShapeSize(dWkvShape); int64_t dWgateSize GetShapeSize(dWgateShape); int64_t dApeSize GetShapeSize(dApeShape); std::vectorbfloat16 xHostData(xSize, bfloat16(0.1f)); std::vectorbfloat16 wkvHostData(wkvSize, bfloat16(0.1f)); std::vectorbfloat16 wgateHostData(wgateSize, bfloat16(0.1f)); std::vectorbfloat16 dCmpKvHostData(dCmpKvSize, bfloat16(0.1f)); std::vectorfloat_t softmaxScoreHostData(softmaxScoreSize, 0.1f); std::vectorfloat_t kvHostData(kvSize, 0.1f); std::vectorint32_t startPosHostData(B, 0); std::vectorbfloat16 dXHostData(dXSize, bfloat16(0.0f)); std::vectorbfloat16 dWkvHostData(dWkvSize, bfloat16(0.0f)); std::vectorbfloat16 dWgateHostData(dWgateSize, bfloat16(0.0f)); std::vectorfloat_t dApeHostData(dApeSize, 0.0f); // 4. 创建 aclTensor void *xDeviceAddr nullptr; void *wkvDeviceAddr nullptr; void *wgateDeviceAddr nullptr; void *dCmpKvDeviceAddr nullptr; void *softmaxScoreDeviceAddr nullptr; void *kvDeviceAddr nullptr; void *startPosDeviceAddr nullptr; void *dXDeviceAddr nullptr; void *dWkvDeviceAddr nullptr; void *dWgateDeviceAddr nullptr; void *dApeDeviceAddr nullptr; aclTensor *x nullptr; aclTensor *wkv nullptr; aclTensor *wgate nullptr; aclTensor *dCmpKv nullptr; aclTensor *softmaxScore nullptr; aclTensor *kv nullptr; aclTensor *startPos nullptr; aclTensor *dX nullptr; aclTensor *dWkv nullptr; aclTensor *dWgate nullptr; aclTensor *dApe nullptr; ret CreateAclTensor(xHostData, xShape, xDeviceAddr, aclDataType::ACL_BF16, x); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(wkvHostData, wkvShape, wkvDeviceAddr, aclDataType::ACL_BF16, wkv); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(wgateHostData, wgateShape, wgateDeviceAddr, aclDataType::ACL_BF16, wgate); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(dCmpKvHostData, dCmpKvShape, dCmpKvDeviceAddr, aclDataType::ACL_BF16, dCmpKv); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(softmaxScoreHostData, softmaxScoreShape, softmaxScoreDeviceAddr, aclDataType::ACL_FLOAT, softmaxScore); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(kvHostData, kvShape, kvDeviceAddr, aclDataType::ACL_FLOAT, kv); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(startPosHostData, startPosShape, startPosDeviceAddr, aclDataType::ACL_INT32, startPos); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(dXHostData, dXShape, dXDeviceAddr, aclDataType::ACL_BF16, dX); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(dWkvHostData, dWkvShape, dWkvDeviceAddr, aclDataType::ACL_BF16, dWkv); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(dWgateHostData, dWgateShape, dWgateDeviceAddr, aclDataType::ACL_BF16, dWgate); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(dApeHostData, dApeShape, dApeDeviceAddr, aclDataType::ACL_FLOAT, dApe); CHECK_RET(ret ACL_SUCCESS, return ret); // 5. 调用 aclnnCompressorGradGetWorkspaceSize uint64_t workspaceSize 0; aclOpExecutor *executor nullptr; ret aclnnCompressorGradGetWorkspaceSize(x, wkv, wgate, dCmpKv, softmaxScore, kv, nullptr, nullptr, startPos, cmpRatio, coff, dX, dWkv, dWgate, dApe, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnCompressorGradGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 6. 申请 workspace void *workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 7. 调用 aclnnCompressorGrad ret aclnnCompressorGrad(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnCompressorGrad failed. ERROR: %d\n, ret); return ret); // 8. 同步等待 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 9. 获取输出 LOG_PRINT(CompressorGrad execution succeeded.\n); PrintBf16Result(dXShape, dXDeviceAddr); PrintBf16Result(dWkvShape, dWkvDeviceAddr); PrintBf16Result(dWgateShape, dWgateDeviceAddr); PrintF32Result(dApeShape, dApeDeviceAddr); // 10. 释放资源 aclDestroyTensor(x); aclDestroyTensor(wkv); aclDestroyTensor(wgate); aclDestroyTensor(dCmpKv); aclDestroyTensor(softmaxScore); aclDestroyTensor(kv); aclDestroyTensor(startPos); aclDestroyTensor(dX); aclDestroyTensor(dWkv); aclDestroyTensor(dWgate); aclDestroyTensor(dApe); aclrtFree(xDeviceAddr); aclrtFree(wkvDeviceAddr); aclrtFree(wgateDeviceAddr); aclrtFree(dCmpKvDeviceAddr); aclrtFree(softmaxScoreDeviceAddr); aclrtFree(kvDeviceAddr); aclrtFree(startPosDeviceAddr); aclrtFree(dXDeviceAddr); aclrtFree(dWkvDeviceAddr); aclrtFree(dWgateDeviceAddr); aclrtFree(dApeDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtDestroyContext(context); aclrtResetDevice(deviceId); aclFinalize(); return 0; }调用示例要点解读示例演示的是C4A 场景D512、coff2、cmp_ratio4、BSH 布局、BF16其完整流程即标准两段式调用aclInit/aclrtSetDevice/aclrtCreateContext/aclrtCreateStream完成 device 与 stream 初始化依据接口 shape 约束构造 10 个 Tensor 的 shape 与 host 数据并通过aclrtMallocaclrtMemcpy创建 Device 侧aclTensor中间张量softmaxScore/kv使用ACL_FLOAT输入与dX/dWkv/dWgate使用ACL_BF16startPos使用ACL_INT32dApe使用ACL_FLOAT与参数表的 dtype 完全对应先调用aclnnCompressorGradGetWorkspaceSize获得workspaceSize与executor按workspaceSize申请 workspace再调用aclnnCompressorGrad执行aclrtSynchronizeStream同步后打印结果最后依次释放 Tensor、内存、stream、context 并aclFinalize。示例中 BSH 布局下cuSeqlens、seqused均传nullptr仅传入startPosshape [B]与约束说明中BS非合轴时 cuSeqlens 必须为空的规定一致。仓库源码级实现解析算子定义与数据类型约束compressor_grad_def.cpp 通过OpDef注册算子6 个必选输入x、wkv、wgate、d_cmp_kv、softmax_score、kv其中x/wkv/wgate/d_cmp_kv支持DT_BF16/DT_FLOAT16softmax_score/kv仅支持DT_FLOATFLOAT323 个可选输入cu_seqlens、seqused、start_pos均要求DT_INT324 个输出d_x、d_wkv、d_wgateBF16/FP16与d_apeFLOAT322 个属性cmp_ratio必选默认值 4与coff可选默认值 1源码常量CMP_RATIO_VALUE 4、COFF_VALUE 1与文档一致所有输入均声明AutoContiguous()且仅支持FORMAT_ND与参数表中非连续Tensor×的约束呼应。Shape 推导compressor_grad_infershape.cpp 实现InferShapeCompressorGrad与InferDataTypeCompressorGrad依据 x 的维度数区分布局3D 为 BS 非合轴B、S、H2D 为 BS 合轴T、H同时校验cmp_ratio必须在 [2, 128] 区间输出推导规则d_x与 x 同 shaped_wkv/d_wgate均为[totalDim, H]totalDim 取 wkv 的 dim0即coff*Dd_ape为[cmp_ratio, totalDim]数据类型推导d_x/d_wkv/d_wgate继承 x 的 dtyped_ape继承softmax_score的 dtype即 FLOAT32。Tiling校验、layout 判定与 workspace 计算compressor_grad_tiling.cpp 是 Host 侧核心逻辑执行一系列校验后产出 tiling 数据Layout 判定x 为 3 维时走LAYOUT_BSH2 维时走LAYOUT_THCheckRequiredInOutExistence强制 TH 布局下cu_seqlens必须非空、BSH 布局下必须为空一致性校验CheckShapeConsistency校验 wkv/wgate 的 dim1 等于 hiddenSize、dim0 等于coff*headDim并分别校验 2D/3D 场景下 dCmpKv、softmaxScore、kv 的各维与batchSize、coff*cmp_ratio、headDim一致CheckDtypeConsistency要求dCmpKv及三个 BF16/FP16 输出与 x 同 dtypeBlock 维约束CheckBlockDimConstrain要求 AIC 核数不小于headDim / D_BASE_SIZED_BASE_SIZE128即最大的 dBaseSizeTilingKey 编码GenTilingKey通过GET_TPL_TILING_KEY(coff, layout, dtype)将Coff(2bit) | Layout(1bit) | DataType(2bit)编码为编译期常量索引Workspace 计算CalcWorkSpace按 FP32 元素数统计五类分区——ape、dXdbRatio2双缓冲、dWkv/dWgate各占一份 dWeightWorkSpaceSize、x 重排缓冲、dXCache最终叠加libapiSize_得到总 workspace 大小并写入context_-workSpaces[0]调度模式使用BATCH_MODE_SCHEDULESyncAll 批模式确保所有核同时启动避免多流方式下执行卡死。Kernel三阶段流水线op_kernel/compressor_grad.py 以 PyPTOpypto_pro实现 NPU kernel其文件头注释给出了统一的三阶段流水线设计Phase 1Vec Scatter每个压缩块计算d_kv、d_score与 APE 累加cast 后插入 L1 并搬运 x。其中_vf_scatter_backward实现 Hadamard 反向与 softmax 反向第一轮计算d_kv dC*sm并用 Dekker 双乘积 单链 Kahan 补偿累加计算列和w Σ(sm*kv)第二轮通过 Knuth TwoSum 将kv - w精确分解为(s, e)对以d_score sm*(dC*s dC*e)的形式重构 softmax 反向结果从而把(dC*kv − col)的大幅相消拆成 O(1)×O(1) 的小项相消Phase 2Cube Matmuldkv wkv dsb wgate → dX partial_compute_dx_partial以及dkv x dsb x → dWkv/dWgate partial_compute_dw_partialweight 梯度首轮覆盖写、后续轮次使用AtomicAdd跨轮累加Phase 3Vec Reduce跨核归约dX/dWkv/dWgate/dApe。其中dApe的归约使用 Kahan 补偿累加_vf_kahan_accumulate因为 d_ape 是纯 Vec 求和、无 matmul 洗噪长序列多项相消场景下舍入噪声会被放大。此外 kernel 还包含_zero_outputs当无压缩块时对四个输出显式刷 0不依赖归约路径的隐式行为。测试与验证仓库在 tests/pytest 下提供了完整的验证体系compressor_grad_cpu_golden.py 实现了 CPU golden 参考按逐 batch、逐压缩块顺序执行反步 1→3Hadamard 反向 → softmax 反向 → 梯度路由映射回 flatnew_kv/new_score再通过 matmul 反向得到d_x/d_wkv/d_wgate/d_ape其推导过程与接口文档的公式一一对应test_compressor_grad_cases.py 定义了覆盖 BSH/TH 两种布局、coff1/2、D128/512、不同 cmp_ratio、seqused/start_pos 组合以及确定性is_deterministic开关的 fuzz 用例compressor_forward.py、compressor_grad_npu.py、compressor_grad_check.py 分别负责前向中间结果构造、NPU 端执行与精度对比Host 侧另有 test_compressor_grad_tiling.cpp 用于 tiling 单测。调用方式小结根据 算子 READMECompressorGrad存在两种调用入口调用方式说明aclnn 接口通过 aclnnCompressorGrad 直接调用 CompressorGrad 算子。PyTorch API通过cann_ops_transformer.compressor调用 Compressor 算子见 torchapi_compressor 文档反向经其 autograd 机制自动调用 CompressorGrad。使用 aclnn 接口时务必遵循两段式流程先aclnnCompressorGradGetWorkspaceSize获取 workspace 大小并完成入参校验再申请 Device 侧 workspace 后调用aclnnCompressorGrad执行最后通过aclrtSynchronizeStream同步并释放全部资源。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考