Star1076
2618
代码
代码
Issues214
Pull Requests496
Actions
项目讨论
Wiki
分析
Star1076
2618
已合并
新增CompressorGrad算子 #9763
莫允扬创建于 19 天前
新增CompressorGrad算子 #9763
已合并
莫允扬创建于 19 天前
从myy268/ops-transformer_cp_gradmaster合入到cann/ops-transformermaster
共 43 个文件变更+9062-190
Mattention/compressor/docs/aclnnCompressor.md+13-6
@@ -101,7 +101,10 @@ aclnnStatus aclnnCompressorGetWorkspaceSize(
101 int64_t coff,101 int64_t coff,
102 int64_t cacheMode,102 int64_t cacheMode,
103 int64_t stateCacheStrideDim0,103 int64_t stateCacheStrideDim0,
104+ bool gradEnabled,
104 const aclTensor *cmpKvOut,105 const aclTensor *cmpKvOut,
106+ const aclTensor *softmaxScoreOut,
107+ const aclTensor *kvOut,
105 uint64_t *workspaceSize,108 uint64_t *workspaceSize,
106 aclOpExecutor **executor)109 aclOpExecutor **executor)
107```110```
@@ -125,20 +128,24 @@ aclnnStatus aclnnCompressor(
125 | wgate | 输入 | 公式中的$W^{Gate}$，表示gate压缩权重。 |不支持空Tensor。| FLOAT16、BFLOAT16 | ND |[coff* D,H]|×|128 | wgate | 输入 | 公式中的$W^{Gate}$，表示gate压缩权重。 |不支持空Tensor。| FLOAT16、BFLOAT16 | ND |[coff* D,H]|×|
126 | stateCacheRef | 输入 | 公式中的$\left[kv\_state, score\_state\right]$, 表示kv\_state和score\_state的历史数据。 |不支持空Tensor| FLOAT32 | ND |[block_num,block_size,2*coff* D]|支持0轴非连续|129 | stateCacheRef | 输入 | 公式中的$\left[kv\_state, score\_state\right]$, 表示kv\_state和score\_state的历史数据。 |不支持空Tensor| FLOAT32 | ND |[block_num,block_size,2*coff* D]|支持0轴非连续|
127 | ape | 输入 | 公式中的$Ape$，表示positional biases。 | 不支持空Tensor。|FLOAT32 | ND |[cmp_ratio,coff* D]|×|130 | ape | 输入 | 公式中的$Ape$，表示positional biases。 | 不支持空Tensor。|FLOAT32 | ND |[cmp_ratio,coff* D]|×|
128- | stateBlockTable | 可选输入 | 表示state\_cache存储使用的block映射表。|当其中元素的值为0时，表示当前位置无需进行更新state\_cache操作；支持S=0,T=0的空Tensor。| INT32 | ND |cache_mode=1时，shape为[B,ceil(Smax/block_size)]，Smax为每个Batch中最大的Sequence Length，当x的shape为[B,S,H]时，Smax=max(start_pos)+S。当x的shape为[T,H]时，Smax=max(start_pos)+max(cu_seqlens[n+1] - cu_seqlens[n])。cache_mode=2时，shape为[B]。当其中元素的值为0时，表示当前位置无需进行更新state_cache操作|×|131+ | stateBlockTableOptional | 可选输入 | 表示state\_cache存储使用的block映射表。|当其中元素的值为0时，表示当前位置无需进行更新state\_cache操作；支持S=0,T=0的空Tensor。| INT32 | ND |cache_mode=1时，shape为[B,ceil(Smax/block_size)]，Smax为每个Batch中最大的Sequence Length，当x的shape为[B,S,H]时，Smax=max(start_pos)+S。当x的shape为[T,H]时，Smax=max(start_pos)+max(cu_seqlens[n+1] - cu_seqlens[n])。cache_mode=2时，shape为[B]。当其中元素的值为0时，表示当前位置无需进行更新state_cache操作|×|
129- | cuSeqlens | 可选输入 | 表示不同Batch中的有效token数。 |支持B=0,S=0,T=0的空Tensor；当x的shape为[B,S,H]时，参数必须为空。| INT32 | ND |当x的shape为[T,H]时，输入shape为[B+1,]|×|132+ | cuSeqlensOptional | 可选输入 | 表示不同Batch中的有效token数。 |支持B=0,S=0,T=0的空Tensor；当x的shape为[B,S,H]时，参数必须为空。| INT32 | ND |当x的shape为[T,H]时，输入shape为[B+1,]|×|
130- | seqused | 可选输入 | 表示不同Batch中实际参与压缩的token数。 |如果指定为None时，表示和每个Batch上的Sequence Length长度相同；支持B=0的空Tensor；如果指定为None时，表示和每个Batch上的Sequence Length长度相同。该入参中每个Batch的有效token数要求小于等于对应Sequence Length长度。当x的shape为[B,S,H]时，要求seqused[n] <= S，且不小于0；当x的shape为[T,H]时，要求seqused[n] <= cu\_seqlens[n+1] - cu\_seqlens[n]，且不小于0。| INT32 | ND |[B,]|×|133+ | sequsedOptional | 可选输入 | 表示不同Batch中实际参与压缩的token数。 |如果指定为None时，表示和每个Batch上的Sequence Length长度相同；支持B=0的空Tensor；如果指定为None时，表示和每个Batch上的Sequence Length长度相同。该入参中每个Batch的有效token数要求小于等于对应Sequence Length长度。当x的shape为[B,S,H]时，要求seqused[n] <= S，且不小于0；当x的shape为[T,H]时，要求seqused[n] <= cu\_seqlens[n+1] - cu\_seqlens[n]，且不小于0。| INT32 | ND |[B,]|×|
131- | startPos | 可选输入 | 表示计算起始位置。 |支持B=0,T=0的空Tensor；当输入为None时，表示从0开始进行计算。| INT32 | ND |[B,]|×|134+ | startPosOptional | 可选输入 | 表示计算起始位置。 |支持B=0,T=0的空Tensor；当输入为None时，表示从0开始进行计算。| INT32 | ND |[B,]|×|
132 | cmpRatio | 输入 | 用于稀疏计算，表示数据压缩率。 |取值范围为[2, 128]内的整数。| INT32 | - |-|-|135 | cmpRatio | 输入 | 用于稀疏计算，表示数据压缩率。 |取值范围为[2, 128]内的整数。| INT32 | - |-|-|
133 | coff | 可选输入 | 表示是否进行overlap数据重排。 |取值范围为[1, 2]。当coff=1时，无需进行overlap数据重排。当coff=2时，需要进行overlap数据重排。| INT32 | - |-|-|136 | coff | 可选输入 | 表示是否进行overlap数据重排。 |取值范围为[1, 2]。当coff=1时，无需进行overlap数据重排。当coff=2时，需要进行overlap数据重排。| INT32 | - |-|-|
134 | cacheMode | 可选输入 | 表示state_cache的存储模式。 |取值范围为[1, 2]；1表示连续buffer，2表示循环buffer。| INT32 | - |-|-|137 | cacheMode | 可选输入 | 表示state_cache的存储模式。 |取值范围为[1, 2]；1表示连续buffer，2表示循环buffer。| INT32 | - |-|-|
135 | stateCacheStrideDim0 | 可选输入 | 表示state_cache的0轴stride。 |-| INT32 | - |-|-|138 | stateCacheStrideDim0 | 可选输入 | 表示state_cache的0轴stride。 |-| INT32 | - |-|-|
136- | cmpKv | 输出 | 表示压缩后的数据。 |支持B=0,S=0,T=0的空Tensor。| FLOAT16、BFLOAT16 | ND |BS合轴：[min(T,T//cmp_ratio+B),D]、BS非合轴：[B,ceil(S/cmp_ratio),D]|×|139+ | gradEnabled | 可选输入 | 表示是否导出softmax\_score/kv中间结果。 | 取值范围为true/false，默认false。当值为true时，算子输出softmax\_score与kv中间结果（供反向传播使用）；值为false时softmax\_score与kv输出内容无效。 | BOOL | - | - | - |
140+ | cmpKvOut | 输出 | 表示压缩后的数据。 |支持B=0,S=0,T=0的空Tensor。| FLOAT16、BFLOAT16 | ND |BS合轴：[min(T,T//cmp_ratio+B),D]、BS非合轴：[B,ceil(S/cmp_ratio),D]|×|
141+ | softmaxScoreOut | 输出 | 公式中的$S^\prime$，表示分组softmax结果。 | 仅在gradEnabled为true时输出有效；支持B=0,S=0,T=0的空Tensor。 | FLOAT32 | ND | BS合轴：[min(T,T//cmp_ratio+B), coff*cmp_ratio, D]、BS非合轴：[B,ceil(S/cmp_ratio),coff*cmp_ratio,D] | × |
142+ | kvOut | 输出 | 公式中的$(S_H)_i$，表示softmax结果与kv\_state的Hadamard乘积。 | 仅在gradEnabled为true时输出有效；支持B=0,S=0,T=0的空Tensor。 | FLOAT32 | ND | 同softmaxScoreOut | × |
137 143 
138<!-- npu="A3" id6 -->144<!-- npu="A3" id6 -->
139- <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：145- <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：
140 cacheMode不支持输入2，且不支持0轴非连续。146 cacheMode不支持输入2，且不支持0轴非连续。
141- cmp_ratio仅支持2/4/8/16/32/64/128147+ cmp_ratio仅支持2/4/8/16/32/64/128。
148+ gradEnabled不支持为true。
142 149 
143<!-- end id6 -->150<!-- end id6 -->
144- **返回值**151- **返回值**
Mattention/compressor/op_host/arch22/compressor_tiling.cpp+1-1文件内容审核中，请稍后刷新重试
Mattention/compressor/op_host/arch35/compressor_tiling.cpp+43-6
@@ -43,6 +43,11 @@ void CompressorTiling::ConvertRequiredParams(gert::TilingContext &context, Compr
43 compressorContext.cmpKv.desc = context.GetOutputDesc(CMP_KV_OUTPUT_INDEX);43 compressorContext.cmpKv.desc = context.GetOutputDesc(CMP_KV_OUTPUT_INDEX);
44 compressorContext.cmpKv.shape = context.GetOutputShape(CMP_KV_OUTPUT_INDEX);44 compressorContext.cmpKv.shape = context.GetOutputShape(CMP_KV_OUTPUT_INDEX);
45 45 
46+ compressorContext.softmaxScore.desc = context.GetOutputDesc(SOFTMAX_SCORE_OUTPUT_INDEX);
47+ compressorContext.softmaxScore.shape = context.GetOutputShape(SOFTMAX_SCORE_OUTPUT_INDEX);
48+ compressorContext.kv.desc = context.GetOutputDesc(KV_OUTPUT_INDEX);
49+ compressorContext.kv.shape = context.GetOutputShape(KV_OUTPUT_INDEX);
50+ 
46 compressorContext.dtype = compressorContext.x.desc->GetDataType();51 compressorContext.dtype = compressorContext.x.desc->GetDataType();
47 auto xDimNum = compressorContext.x.shape->GetStorageShape().GetDimNum();52 auto xDimNum = compressorContext.x.shape->GetStorageShape().GetDimNum();
48 if (xDimNum == COMPRESSOR_DIM_NUM_3) {53 if (xDimNum == COMPRESSOR_DIM_NUM_3) {
@@ -89,6 +94,7 @@ ge::graphStatus CompressorTiling::ConvertContext(gert::TilingContext &context, C
89 compressorContext.stateCacheStrideDim0 = attrs->GetAttrPointer<int>(STATE_CACHE_STRIDE_DIM0_ATTR_INDEX);94 compressorContext.stateCacheStrideDim0 = attrs->GetAttrPointer<int>(STATE_CACHE_STRIDE_DIM0_ATTR_INDEX);
90 compressorContext.batchConsistency = context.GetDeterministicLevel();95 compressorContext.batchConsistency = context.GetDeterministicLevel();
91 OP_LOGD(context.GetNodeName(), "deterministic_level=%d", context.GetDeterministicLevel());96 OP_LOGD(context.GetNodeName(), "deterministic_level=%d", context.GetDeterministicLevel());
97+ compressorContext.gradEnabled = attrs->GetAttrPointer<bool>(GRAD_ENABLED_ATTR_INDEX);
92 OP_CHECK_IF(context.GetWorkspaceSizes(1) == nullptr,98 OP_CHECK_IF(context.GetWorkspaceSizes(1) == nullptr,
93 OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.GetNodeName(), "workSpaceSize",99 OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.GetNodeName(), "workSpaceSize",
94 "got from ge is nullptr"),100 "got from ge is nullptr"),
@@ -352,6 +358,7 @@ ge::graphStatus CompressorTiling::GenTilingKey() const
352 uint8_t layout = 0;358 uint8_t layout = 0;
353 uint8_t templateId = static_cast<uint8_t>(context_->templateId);359 uint8_t templateId = static_cast<uint8_t>(context_->templateId);
354 uint8_t cacheMode = static_cast<uint8_t>(*context_->cacheMode);360 uint8_t cacheMode = static_cast<uint8_t>(*context_->cacheMode);
361+ uint8_t gradEnabled = static_cast<uint8_t>(context_->gradEnabled ? 1 : 0);
355 362 
356 auto xDtype = context_->x.desc->GetDataType();363 auto xDtype = context_->x.desc->GetDataType();
357 if (xDtype == ge::DT_BF16) {364 if (xDtype == ge::DT_BF16) {
@@ -366,9 +373,10 @@ ge::graphStatus CompressorTiling::GenTilingKey() const
366 layout = 1;373 layout = 1;
367 }374 }
368 375 
369- context_->tilingKey = GET_TPL_TILING_KEY(layout, dtype, coff, cacheMode, templateId);376+ context_->tilingKey = GET_TPL_TILING_KEY(layout, dtype, coff, cacheMode, templateId, gradEnabled);
370- OP_LOGI(context_->opName, "Compressor dtype:%hhu layout:%hhu coff:%hhu, cacheMode: %u, template_id:%hhu", dtype,377+ OP_LOGI(context_->opName,
371- layout, coff, cacheMode, templateId);378+ "Compressor dtype:%hhu layout:%hhu coff:%hhu, cacheMode: %u, template_id:%hhu grad_enabled:%hhu", dtype,
379+ layout, coff, cacheMode, templateId, gradEnabled);
372 OP_LOGI(context_->opName, "Compressor tilingKey:%lu", context_->tilingKey);380 OP_LOGI(context_->opName, "Compressor tilingKey:%lu", context_->tilingKey);
373 381 
374 return ge::GRAPH_SUCCESS;382 return ge::GRAPH_SUCCESS;
@@ -382,8 +390,8 @@ ge::graphStatus CompressorTiling::CheckSinglePara() const
382 ge::GRAPH_SUCCESS != CheckSingleParaCuSeqlens() || ge::GRAPH_SUCCESS != CheckSingleParaSeqused() ||390 ge::GRAPH_SUCCESS != CheckSingleParaCuSeqlens() || ge::GRAPH_SUCCESS != CheckSingleParaSeqused() ||
383 ge::GRAPH_SUCCESS != CheckSingleParaStartPos() || ge::GRAPH_SUCCESS != CheckSingleParaCmpKv() ||391 ge::GRAPH_SUCCESS != CheckSingleParaStartPos() || ge::GRAPH_SUCCESS != CheckSingleParaCmpKv() ||
384 ge::GRAPH_SUCCESS != CheckSingleParaCmpRatio() || ge::GRAPH_SUCCESS != CheckSingleParaCoff() ||392 ge::GRAPH_SUCCESS != CheckSingleParaCmpRatio() || ge::GRAPH_SUCCESS != CheckSingleParaCoff() ||
385- ge::GRAPH_SUCCESS != CheckSingleParaCacheMode()) {393+ ge::GRAPH_SUCCESS != CheckSingleParaCacheMode() || ge::GRAPH_SUCCESS != CheckSingleParaSoftmaxScore() ||
386- return ge::GRAPH_FAILED;394+ ge::GRAPH_SUCCESS != CheckSingleParaKv()) { return ge::GRAPH_FAILED;
387 }395 }
388 return ge::GRAPH_SUCCESS;396 return ge::GRAPH_SUCCESS;
389}397}
@@ -635,8 +643,25 @@ ge::graphStatus CompressorTiling::CheckSingleParaCmpKv() const
635 return ge::GRAPH_SUCCESS;643 return ge::GRAPH_SUCCESS;
636}644}
637 645 
638-ge::graphStatus CompressorTiling::CheckSingleParaCmpRatio() const646+ge::graphStatus CompressorTiling::CheckSingleParaSoftmaxScore() const
639{647{
648+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->softmaxScore.desc, SOFTMAX_SCORE_NAME) ||
649+ ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->softmaxScore.shape, SOFTMAX_SCORE_NAME)) {
650+ return ge::GRAPH_FAILED;
651+ }
652+ return ge::GRAPH_SUCCESS;
653+}
654+ 
655+ge::graphStatus CompressorTiling::CheckSingleParaKv() const
656+{
657+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->kv.desc, KV_NAME) ||
658+ ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->kv.shape, KV_NAME)) {
659+ return ge::GRAPH_FAILED;
660+ }
661+ return ge::GRAPH_SUCCESS;
662+}
663+ 
664+ge::graphStatus CompressorTiling::CheckSingleParaCmpRatio() const{
640 uint32_t cmpRatio = static_cast<uint32_t>(*context_->cmpRatio);665 uint32_t cmpRatio = static_cast<uint32_t>(*context_->cmpRatio);
641 OP_CHECK_IF(cmpRatio > MAX_CMPRATIO_SIZE || cmpRatio < MIN_CMPRATIO_SIZE,666 OP_CHECK_IF(cmpRatio > MAX_CMPRATIO_SIZE || cmpRatio < MIN_CMPRATIO_SIZE,
642 OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->opName, "cmp_ratio", std::to_string(cmpRatio),667 OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->opName, "cmp_ratio", std::to_string(cmpRatio),
@@ -714,6 +739,18 @@ ge::graphStatus CompressorTiling::CheckRequiredInOutExistence() const
714 OP_CHECK_IF(context_->cmpKv.desc == nullptr,739 OP_CHECK_IF(context_->cmpKv.desc == nullptr,
715 OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "cmp_kv", "desc is nullptr"),740 OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "cmp_kv", "desc is nullptr"),
716 return ge::GRAPH_FAILED);741 return ge::GRAPH_FAILED);
742+ OP_CHECK_IF(context_->softmaxScore.shape == nullptr,
743+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "softmax_score", "shape is nullptr"),
744+ return ge::GRAPH_FAILED);
745+ OP_CHECK_IF(context_->softmaxScore.desc == nullptr,
746+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "softmax_score", "desc is nullptr"),
747+ return ge::GRAPH_FAILED);
748+ OP_CHECK_IF(context_->kv.shape == nullptr,
749+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "kv", "shape is nullptr"),
750+ return ge::GRAPH_FAILED);
751+ OP_CHECK_IF(context_->kv.desc == nullptr,
752+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "kv", "desc is nullptr"),
753+ return ge::GRAPH_FAILED);
717 if (context_->layout == LayoutType::LAYOUT_TH) {754 if (context_->layout == LayoutType::LAYOUT_TH) {
718 OP_CHECK_IF(context_->cuSeqlens.desc == nullptr,755 OP_CHECK_IF(context_->cuSeqlens.desc == nullptr,
719 OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "cu_seqlens",756 OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "cu_seqlens",
Mattention/compressor/op_host/arch35/compressor_tiling.h+36-20
@@ -57,9 +57,12 @@ constexpr uint32_t CMP_RATIO_ATTR_INDEX = 0;
57constexpr uint32_t COFF_ATTR_INDEX = 1;57constexpr uint32_t COFF_ATTR_INDEX = 1;
58constexpr uint32_t CACHE_MODE_ATTR_INDEX = 2;58constexpr uint32_t CACHE_MODE_ATTR_INDEX = 2;
59constexpr uint32_t STATE_CACHE_STRIDE_DIM0_ATTR_INDEX = 3;59constexpr uint32_t STATE_CACHE_STRIDE_DIM0_ATTR_INDEX = 3;
60+constexpr uint32_t GRAD_ENABLED_ATTR_INDEX = 4;
60 61 
61// OUTPUT62// OUTPUT
62constexpr uint32_t CMP_KV_OUTPUT_INDEX = 0;63constexpr uint32_t CMP_KV_OUTPUT_INDEX = 0;
64+constexpr uint32_t SOFTMAX_SCORE_OUTPUT_INDEX = 2;
65+constexpr uint32_t KV_OUTPUT_INDEX = 3;
63 66 
64constexpr uint32_t COMPRESSOR_DIM_NUM_1 = 1;67constexpr uint32_t COMPRESSOR_DIM_NUM_1 = 1;
65constexpr uint32_t COMPRESSOR_DIM_NUM_2 = 2;68constexpr uint32_t COMPRESSOR_DIM_NUM_2 = 2;
@@ -94,32 +97,40 @@ static const std::string CMP_RATIO_NAME = "cmp_ratio";
94static const std::string COFF_NAME = "coff";97static const std::string COFF_NAME = "coff";
95static const std::string CACHE_MODE_NAME = "cache_mode";98static const std::string CACHE_MODE_NAME = "cache_mode";
96static const std::string CMP_KV_NAME = "cmp_kv";99static const std::string CMP_KV_NAME = "cmp_kv";
100+static const std::string SOFTMAX_SCORE_NAME = "softmax_score";
101+static const std::string KV_NAME = "kv";
97 102 
98static std::string DataTypeToSerialString(ge::DataType type);103static std::string DataTypeToSerialString(ge::DataType type);
99 104 
100const std::map<std::string, std::vector<ge::DataType>> DTYPE_SUPPORT_MAP = {105const std::map<std::string, std::vector<ge::DataType>> DTYPE_SUPPORT_MAP = {
101- {X_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},106+ {X_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
102- {WKV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},107+ {WKV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
103- {WGATE_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},108+ {WGATE_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
104- {STATE_CACHE_NAME, {ge::DT_FLOAT}},109+ {STATE_CACHE_NAME, {ge::DT_FLOAT}},
105- {APE_NAME, {ge::DT_FLOAT}},110+ {APE_NAME, {ge::DT_FLOAT}},
106- {STATE_BLOCK_TABLE_NAME, {ge::DT_INT32}},111+ {STATE_BLOCK_TABLE_NAME, {ge::DT_INT32}},
107- {CU_SEQLENS_NAME, {ge::DT_INT32}},112+ {CU_SEQLENS_NAME, {ge::DT_INT32}},
108- {SEQUSED_NAME, {ge::DT_INT32}},113+ {SEQUSED_NAME, {ge::DT_INT32}},
109- {START_POS_NAME, {ge::DT_INT32}},114+ {START_POS_NAME, {ge::DT_INT32}},
110- {CMP_KV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}}};115+ {CMP_KV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
116+ {SOFTMAX_SCORE_NAME, {ge::DT_FLOAT}},
117+ {KV_NAME, {ge::DT_FLOAT}}
118+};
111 119 
112const std::map<std::string, std::vector<uint32_t>> DIM_NUM_MAP = {120const std::map<std::string, std::vector<uint32_t>> DIM_NUM_MAP = {
113- {X_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}},121+ {X_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}},
114- {WKV_NAME, {COMPRESSOR_DIM_NUM_2}},122+ {WKV_NAME, {COMPRESSOR_DIM_NUM_2}},
115- {WGATE_NAME, {COMPRESSOR_DIM_NUM_2}},123+ {WGATE_NAME, {COMPRESSOR_DIM_NUM_2}},
116- {STATE_CACHE_NAME, {COMPRESSOR_DIM_NUM_3}},124+ {STATE_CACHE_NAME, {COMPRESSOR_DIM_NUM_3}},
117- {APE_NAME, {COMPRESSOR_DIM_NUM_2}},125+ {APE_NAME, {COMPRESSOR_DIM_NUM_2}},
118- {STATE_BLOCK_TABLE_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_1}},126+ {STATE_BLOCK_TABLE_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_1}},
119- {CU_SEQLENS_NAME, {COMPRESSOR_DIM_NUM_1}},127+ {CU_SEQLENS_NAME, {COMPRESSOR_DIM_NUM_1}},
120- {SEQUSED_NAME, {COMPRESSOR_DIM_NUM_1}},128+ {SEQUSED_NAME, {COMPRESSOR_DIM_NUM_1}},
121- {START_POS_NAME, {COMPRESSOR_DIM_NUM_1}},129+ {START_POS_NAME, {COMPRESSOR_DIM_NUM_1}},
122- {CMP_KV_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}}};130+ {CMP_KV_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}},
131+ {SOFTMAX_SCORE_NAME, {COMPRESSOR_DIM_NUM_3, COMPRESSOR_DIM_NUM_4}},
132+ {KV_NAME, {COMPRESSOR_DIM_NUM_3, COMPRESSOR_DIM_NUM_4}}
133+};
123 134 
124static const std::map<std::string, uint32_t> LAYOUT_DIM_MAP = {135static const std::map<std::string, uint32_t> LAYOUT_DIM_MAP = {
125 {"BSH", COMPRESSOR_DIM_NUM_3},136 {"BSH", COMPRESSOR_DIM_NUM_3},
@@ -216,11 +227,14 @@ struct CompressorContext {
216 OptionalParaInfo seqUsed;227 OptionalParaInfo seqUsed;
217 OptionalParaInfo startPos;228 OptionalParaInfo startPos;
218 RequiredParaInfo cmpKv;229 RequiredParaInfo cmpKv;
230+ RequiredParaInfo softmaxScore;
231+ RequiredParaInfo kv;
219 232 
220 const int *coff;233 const int *coff;
221 const int *cmpRatio;234 const int *cmpRatio;
222 const int *cacheMode;235 const int *cacheMode;
223 const int *stateCacheStrideDim0;236 const int *stateCacheStrideDim0;
237+ const bool *gradEnabled;
224 TemplateId templateId;238 TemplateId templateId;
225 239 
226 ge::DataType dtype = ge::DT_BF16;240 ge::DataType dtype = ge::DT_BF16;
@@ -281,6 +295,8 @@ private:
281 ge::graphStatus CheckSingleParaSeqused() const;295 ge::graphStatus CheckSingleParaSeqused() const;
282 ge::graphStatus CheckSingleParaStartPos() const;296 ge::graphStatus CheckSingleParaStartPos() const;
283 ge::graphStatus CheckSingleParaCmpKv() const;297 ge::graphStatus CheckSingleParaCmpKv() const;
298+ ge::graphStatus CheckSingleParaSoftmaxScore() const;
299+ ge::graphStatus CheckSingleParaKv() const;
284 ge::graphStatus CheckSingleParaCmpRatio() const;300 ge::graphStatus CheckSingleParaCmpRatio() const;
285 ge::graphStatus CheckSingleParaCoff() const;301 ge::graphStatus CheckSingleParaCoff() const;
286 ge::graphStatus CheckSingleParaCacheMode() const;302 ge::graphStatus CheckSingleParaCacheMode() const;
Mattention/compressor/op_host/compressor_def.cpp+9-0
@@ -75,10 +75,19 @@ public:
75 .ParamType(REQUIRED)75 .ParamType(REQUIRED)
76 .DataTypeList({ge::DT_FLOAT})76 .DataTypeList({ge::DT_FLOAT})
77 .FormatList({ge::FORMAT_ND});77 .FormatList({ge::FORMAT_ND});
78+ this->Output("softmax_score")
79+ .ParamType(REQUIRED)
80+ .DataTypeList({ge::DT_FLOAT})
81+ .FormatList({ge::FORMAT_ND});
82+ this->Output("kv")
83+ .ParamType(REQUIRED)
84+ .DataTypeList({ge::DT_FLOAT})
85+ .FormatList({ge::FORMAT_ND});
78 this->Attr("cmp_ratio").AttrType(REQUIRED).Int(CMP_RATIO_VALUE);86 this->Attr("cmp_ratio").AttrType(REQUIRED).Int(CMP_RATIO_VALUE);
79 this->Attr("coff").AttrType(OPTIONAL).Int(COFF_VALUE);87 this->Attr("coff").AttrType(OPTIONAL).Int(COFF_VALUE);
80 this->Attr("cache_mode").AttrType(OPTIONAL).Int(CACHE_MODE_VALUE);88 this->Attr("cache_mode").AttrType(OPTIONAL).Int(CACHE_MODE_VALUE);
81 this->Attr("state_cache_stride_dim0").AttrType(OPTIONAL).Int(STATE_CACHE_STRIDE_DIM0);89 this->Attr("state_cache_stride_dim0").AttrType(OPTIONAL).Int(STATE_CACHE_STRIDE_DIM0);
90+ this->Attr("grad_enabled").AttrType(OPTIONAL).Bool(false);
82 OpAICoreConfig aicore_config;91 OpAICoreConfig aicore_config;
83 aicore_config.DynamicCompileStaticFlag(true)92 aicore_config.DynamicCompileStaticFlag(true)
84 .DynamicFormatFlag(true)93 .DynamicFormatFlag(true)
Mattention/compressor/op_host/compressor_infershape.cpp+40-0
@@ -39,6 +39,8 @@ namespace ops {
39 39 
40 // OUTPUT40 // OUTPUT
41 constexpr uint32_t CMP_KV_OUTPUT_INDEX = 0;41 constexpr uint32_t CMP_KV_OUTPUT_INDEX = 0;
42+ constexpr uint32_t SOFTMAX_SCORE_OUTPUT_INDEX = 2;
43+ constexpr uint32_t KV_OUTPUT_INDEX = 3;
42 44 
43 // ATTR DEFAULT VALUE45 // ATTR DEFAULT VALUE
44 constexpr uint32_t CMP_RATIO_VALUE = 4;46 constexpr uint32_t CMP_RATIO_VALUE = 4;
@@ -126,16 +128,52 @@ ge::graphStatus SetCompressorShapeDim(const CompressorProtoShapeParam &shapePara
126 OP_CHECK_NULL_WITH_CONTEXT(context, cmpKvShape);128 OP_CHECK_NULL_WITH_CONTEXT(context, cmpKvShape);
127 auto attr = context->GetAttrs();129 auto attr = context->GetAttrs();
128 OP_CHECK_NULL_WITH_CONTEXT(context, attr);130 OP_CHECK_NULL_WITH_CONTEXT(context, attr);
131+ const int64_t *cmpRatioPtr = attr->GetAttrPointer<int64_t>(CMP_RATIO_ATTR_INDEX);
132+ int64_t cmpRatio = (cmpRatioPtr != nullptr) ? *cmpRatioPtr : CMP_RATIO_VALUE;
133+ const int64_t *coffPtr = attr->GetAttrPointer<int64_t>(COFF_ATTR_INDEX);
134+ int64_t coff = (coffPtr != nullptr) ? *coffPtr : COFF_VALUE;
135+ int64_t extraDim = coff * cmpRatio;
136+ 
129 // Set output shape137 // Set output shape
130 if (!shapeParam.isBsMerge) {138 if (!shapeParam.isBsMerge) {
131 cmpKvShape->SetDimNum(DIM_NUM_3); // (B, Sr, H)139 cmpKvShape->SetDimNum(DIM_NUM_3); // (B, Sr, H)
132 cmpKvShape->SetDim(DIM_INDEX_0, shapeParam.B);140 cmpKvShape->SetDim(DIM_INDEX_0, shapeParam.B);
133 cmpKvShape->SetDim(DIM_INDEX_1, shapeParam.Sr);141 cmpKvShape->SetDim(DIM_INDEX_1, shapeParam.Sr);
134 cmpKvShape->SetDim(DIM_INDEX_2, shapeParam.D);142 cmpKvShape->SetDim(DIM_INDEX_2, shapeParam.D);
143+ 
144+ auto softmaxScoreShape = context->GetOutputShape(SOFTMAX_SCORE_OUTPUT_INDEX);
145+ OP_CHECK_NULL_WITH_CONTEXT(context, softmaxScoreShape);
146+ softmaxScoreShape->SetDimNum(DIM_NUM_4); // (B, Sr, coff*cmpRatio, D)
147+ softmaxScoreShape->SetDim(DIM_INDEX_0, shapeParam.B);
148+ softmaxScoreShape->SetDim(DIM_INDEX_1, shapeParam.Sr);
149+ softmaxScoreShape->SetDim(DIM_INDEX_2, extraDim);
150+ softmaxScoreShape->SetDim(DIM_INDEX_3, shapeParam.D);
151+ 
152+ auto kvShape = context->GetOutputShape(KV_OUTPUT_INDEX);
153+ OP_CHECK_NULL_WITH_CONTEXT(context, kvShape);
154+ kvShape->SetDimNum(DIM_NUM_4); // (B, Sr, coff*cmpRatio, D)
155+ kvShape->SetDim(DIM_INDEX_0, shapeParam.B);
156+ kvShape->SetDim(DIM_INDEX_1, shapeParam.Sr);
157+ kvShape->SetDim(DIM_INDEX_2, extraDim);
158+ kvShape->SetDim(DIM_INDEX_3, shapeParam.D);
135 } else {159 } else {
136 cmpKvShape->SetDimNum(DIM_NUM_2); // (T, N, Hckv)160 cmpKvShape->SetDimNum(DIM_NUM_2); // (T, N, Hckv)
137 cmpKvShape->SetDim(DIM_INDEX_0, shapeParam.Sr);161 cmpKvShape->SetDim(DIM_INDEX_0, shapeParam.Sr);
138 cmpKvShape->SetDim(DIM_INDEX_1, shapeParam.D);162 cmpKvShape->SetDim(DIM_INDEX_1, shapeParam.D);
163+ 
164+ auto softmaxScoreShape = context->GetOutputShape(SOFTMAX_SCORE_OUTPUT_INDEX);
165+ OP_CHECK_NULL_WITH_CONTEXT(context, softmaxScoreShape);
166+ softmaxScoreShape->SetDimNum(DIM_NUM_3); // (Sr, coff*cmpRatio, D)
167+ softmaxScoreShape->SetDim(DIM_INDEX_0, shapeParam.Sr);
168+ softmaxScoreShape->SetDim(DIM_INDEX_1, extraDim);
169+ softmaxScoreShape->SetDim(DIM_INDEX_2, shapeParam.D);
170+ 
171+ auto kvShape = context->GetOutputShape(KV_OUTPUT_INDEX);
172+ OP_CHECK_NULL_WITH_CONTEXT(context, kvShape);
173+ kvShape->SetDimNum(DIM_NUM_3); // (Sr, coff*cmpRatio, D)
174+ kvShape->SetDim(DIM_INDEX_0, shapeParam.Sr);
175+ kvShape->SetDim(DIM_INDEX_1, extraDim);
176+ kvShape->SetDim(DIM_INDEX_2, shapeParam.D);
139 }177 }
140 178 
141 return GRAPH_SUCCESS;179 return GRAPH_SUCCESS;
@@ -149,6 +187,8 @@ ge::graphStatus InferDataTypeCompressor(gert::InferDataTypeContext* context)
149 OP_LOGI(context->GetNodeName(), "Enter Compressor inferDataType impl.");187 OP_LOGI(context->GetNodeName(), "Enter Compressor inferDataType impl.");
150 188 
151 context->SetOutputDataType(CMP_KV_OUTPUT_INDEX, context->GetRequiredInputDataType(TOKEN_X_INPUT_INDEX));189 context->SetOutputDataType(CMP_KV_OUTPUT_INDEX, context->GetRequiredInputDataType(TOKEN_X_INPUT_INDEX));
190+ context->SetOutputDataType(SOFTMAX_SCORE_OUTPUT_INDEX, ge::DT_FLOAT);
191+ context->SetOutputDataType(KV_OUTPUT_INDEX, ge::DT_FLOAT);
152 192 
153 return GRAPH_SUCCESS;193 return GRAPH_SUCCESS;
154}194}
Mattention/compressor/op_kernel/arch22/compressor_block_vec_full_load.h+5-0
@@ -982,6 +982,11 @@ __aicore__ inline void CompressorBlockVectorFullLoad<COMP>::CalcGroupInfo(Vec1Sp
982 if (constInfo_.kBaseNum > 1) {982 if (constInfo_.kBaseNum > 1) {
983 splitInfo.dBaseSize = max(splitInfo.dBaseSize, FP32_REPEAT_ELEMENT_NUM);983 splitInfo.dBaseSize = max(splitInfo.dBaseSize, FP32_REPEAT_ELEMENT_NUM);
984 }984 }
985+ // 32B(8个FP32)对齐的UB列窗口上限（同 arch35 修复）：防止 CalcTilingStrategy
986+ // 拆出非32B对齐的列窗口（DataCopy blockLen/srcGap 错位）。
987+ // cmpRatio<=4 分支无需限制（maxDealColNum >= 2K >= headDim，不触发拆分）。
988+ uint32_t maxDealColNum = BUFFER_SIZE_BYTE_32K / (cmpRatio_ * coff_ * sizeof(T));
989+ splitInfo.dBaseSize = min(splitInfo.dBaseSize, FloorPow2(Trunc(maxDealColNum, BlockElementNum<T>())));
985 }990 }
986 // 结果输出到GM前必须转换成X_T，dBaseSize * sizeof(X_T)需32B对齐991 // 结果输出到GM前必须转换成X_T，dBaseSize * sizeof(X_T)需32B对齐
987 splitInfo.dBaseSize = max(splitInfo.dBaseSize, BlockElementNum<X_T>());992 splitInfo.dBaseSize = max(splitInfo.dBaseSize, BlockElementNum<X_T>());
Mattention/compressor/op_kernel/arch22/compressor_block_vec_perf.h+5-0
@@ -914,6 +914,11 @@ __aicore__ inline void CompressorBlockVectorPerf<COMP>::CalcGroupInfo(const Vec1
914{914{
915 uint32_t aiCoreNum = constInfo_.usedCoreNum * 2;915 uint32_t aiCoreNum = constInfo_.usedCoreNum * 2;
916 splitInfo.dBaseSize = constInfo_.headDim / min(FloorPow2(aiCoreNum), CeilPow2(CeilDivT(aiCoreNum, info.dealTcNum)));916 splitInfo.dBaseSize = constInfo_.headDim / min(FloorPow2(aiCoreNum), CeilPow2(CeilDivT(aiCoreNum, info.dealTcNum)));
917+ // 32B(8个FP32)对齐的UB列窗口上限（同 arch35 NORMAL 修复）：dBaseSize 超过它时
918+ // CalcTilingStrategy 的 dSplitSize = dBaseSize/dLoopCount 整数除法会切出非32B
919+ // 对齐的列窗口（DataCopy blockLen/srcGap 整数除法错位 → 数据错乱）。
920+ uint32_t maxDealColNum = BUFFER_SIZE_BYTE_32K / (cmpRatio_ * coff_ * sizeof(T));
921+ splitInfo.dBaseSize = min(splitInfo.dBaseSize, FloorPow2(Trunc(maxDealColNum, BlockElementNum<T>())));
917 // 结果输出到GM前必须转换成X_T，dBaseSize * sizeof(X_T)需32B对齐922 // 结果输出到GM前必须转换成X_T，dBaseSize * sizeof(X_T)需32B对齐
918 splitInfo.dBaseSize = max(splitInfo.dBaseSize, BlockElementNum<X_T>());923 splitInfo.dBaseSize = max(splitInfo.dBaseSize, BlockElementNum<X_T>());
919 splitInfo.vec1GroupSize = constInfo_.headDim / splitInfo.dBaseSize;924 splitInfo.vec1GroupSize = constInfo_.headDim / splitInfo.dBaseSize;
Mattention/compressor/op_kernel/arch22/compressor_template_tiling_key.h+5-1
@@ -34,7 +34,10 @@ ASCENDC_TPL_ARGS_DECL(compressor, // 算子唯一标识，与opType保持一致
34 // bit:7-8 cache_mode 1:LINEAR_BUFFER 2:RING_BUFFER34 // bit:7-8 cache_mode 1:LINEAR_BUFFER 2:RING_BUFFER
35 ASCENDC_TPL_UINT_DECL(CACHE_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2),35 ASCENDC_TPL_UINT_DECL(CACHE_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2),
36 // bit:9-10 template_id 0:empty_tensor 1:normal 2:full load36 // bit:9-10 template_id 0:empty_tensor 1:normal 2:full load
37- ASCENDC_TPL_UINT_DECL(TEMPLATE_ID, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 0, 1, 2), );37+ ASCENDC_TPL_UINT_DECL(TEMPLATE_ID, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 0, 1, 2),
38+ // bit:11 grad_enabled 0:关闭 1:开启 (arch22仅支持0)
39+ ASCENDC_TPL_BOOL_DECL(GRAD_ENABLED, 0, 1),);
40+ 
38 41 
39ASCENDC_TPL_SEL(42ASCENDC_TPL_SEL(
40 43 
@@ -43,6 +46,7 @@ ASCENDC_TPL_SEL(
43 ASCENDC_TPL_UINT_SEL(COFF, ASCENDC_TPL_UI_LIST, 1, 2),46 ASCENDC_TPL_UINT_SEL(COFF, ASCENDC_TPL_UI_LIST, 1, 2),
44 ASCENDC_TPL_UINT_SEL(CACHE_MODE, ASCENDC_TPL_UI_LIST, 1, 2),47 ASCENDC_TPL_UINT_SEL(CACHE_MODE, ASCENDC_TPL_UI_LIST, 1, 2),
45 ASCENDC_TPL_UINT_SEL(TEMPLATE_ID, ASCENDC_TPL_UI_LIST, 0, 1, 2),48 ASCENDC_TPL_UINT_SEL(TEMPLATE_ID, ASCENDC_TPL_UI_LIST, 0, 1, 2),
49+ ASCENDC_TPL_BOOL_SEL(GRAD_ENABLED, 0),
46 ASCENDC_TPL_TILING_STRUCT_SEL(optiling::CompressorTilingData)), );50 ASCENDC_TPL_TILING_STRUCT_SEL(optiling::CompressorTilingData)), );
47 51 
48#endif // COMPRESSOR_TEMPLATE_TILING_KEY_H52#endif // COMPRESSOR_TEMPLATE_TILING_KEY_H
Mattention/compressor/op_kernel/arch22/compressor_tools.h+80-27
@@ -651,41 +651,91 @@ __aicore__ inline void CompressorVec1SliceIterator<COMP>::IteratorSlice()
651 sliceInfo_.sIdx += sliceInfo_.validSeqCnt;651 sliceInfo_.sIdx += sliceInfo_.validSeqCnt;
652 if (sliceInfo_.sIdx >= sliceInfo_.bSeqUsed) {652 if (sliceInfo_.sIdx >= sliceInfo_.bSeqUsed) {
653 do {653 do {
654- uint32_t seqLength = tools_.GetSeqLength(sliceInfo_.bIdx);654+ const uint32_t seqLength = tools_.GetSeqLength(sliceInfo_.bIdx);
655 if (sliceInfo_.bSeqUsed < seqLength) {655 if (sliceInfo_.bSeqUsed < seqLength) {
656- uint64_t nextAlignSIdx = Align(sliceInfo_.bStartPos + sliceInfo_.sIdx, (uint64_t)cmpRatio) -656+ // ── (A) tailH 行推进（不消耗任务量）──
657- sliceInfo_.bStartPos;657+ // slice 处理到 su 为止，su 到下一个压缩块分界点（全局 cr 对齐的"切分点"）
658- if (nextAlignSIdx >= seqLength) {658+ // 之间的行是 slice 尾部 padding（tailH）。这些行在任务量上已由 slice 的
659- sliceInfo_.dealedSeqCnt += seqLength - sliceInfo_.sIdx;659+ // dealTcSize 消耗，但 workspace 行号（dealedSeqCnt）尚未推进——必须在此
660- sliceInfo_.sIdx = seqLength;660+ // 推进，否则核起点落在非切分点，后续 slice 从错位位置读数据。
661- } else {661+ //
662- sliceInfo_.dealedSeqCnt += nextAlignSIdx - sliceInfo_.sIdx;662+ // ★特殊1：仅 sIdx > 0（slice 确实处理过）时才推进 tailH。
663- uint32_t tcGap = CeilDivT(static_cast<int32_t>(seqLength - nextAlignSIdx),663+ // sIdx == 0 表示 seqused == 0 的无效 batch（无 slice），其行全部属于
664- static_cast<int32_t>(cmpRatio));664+ // 空洞（块任务量未消耗）；此时按 tailH 推进会"行超前于任务量"
665- if (sliceInfo_.bSeqUsed == 0 && nextAlignSIdx > sliceInfo_.sIdx) {665+ // （b=53: 全局对齐量 62 行，但 x 行仅 1 行）→ workspace 偏移虚增越界。
666- // 此时bseqused所在压缩块未被纳入计算666+ //
667- tcGap++;667+ // ★特殊2：对齐量 clamp 到空洞内。对齐量（< cmpRatio）可能超过本 batch
668- }668+ // 的 x 行剩余（seqLength - sIdx，如无效 batch 的 x 行很短），
669- sliceInfo_.sIdx = nextAlignSIdx;669+ // 超出的行属于下一个 batch 或窗口外，不得推进。
670- if (needDealTcSize_ < tcGap) {670+ if (sliceInfo_.sIdx > 0) {
671- sliceInfo_.dealedSeqCnt += needDealTcSize_ * cmpRatio;671+ uint64_t nextAlignSIdx = Align(sliceInfo_.bStartPos + sliceInfo_.sIdx,
672- sliceInfo_.sIdx += needDealTcSize_ * cmpRatio;672+ static_cast<uint64_t>(cmpRatio)) -
673- needDealTcSize_ = 0;673+ sliceInfo_.bStartPos;
674- break;674+ uint32_t align = min(static_cast<uint32_t>(nextAlignSIdx - sliceInfo_.sIdx),
675- }675+ seqLength - sliceInfo_.sIdx);
676- sliceInfo_.dealedSeqCnt += seqLength - sliceInfo_.sIdx;676+ sliceInfo_.dealedSeqCnt += align;
677- needDealTcSize_ -= tcGap;677+ sliceInfo_.sIdx += align;
678 }678 }
679+ 
680+ // ── 空洞（tailH 之后）对应的压缩块数 tcGap（需消耗的任务量）──
681+ // 空洞 = [sIdx, seqLength) 的 x 行，全局位置 [bStartPos+sIdx, bStartPos+seqLength)。
682+ // 块数 = 空洞覆盖的全局 cr 块数。
683+ //
684+ // ★特殊3：sIdx == 0（seqused == 0）时起点块（含头部 padding 的块）
685+ // 未被 slice 消耗，块数从 floor(bStartPos/cr) 起算（= tcNum 公式，起点块计入）；
686+ // sIdx > 0（tailH 已推进到切分点）时从 ceil((bStartPos+sIdx)/cr) 起算
687+ // （起点块已由 slice 的 dealTcSize 消耗）。
688+ const uint32_t gapRows = seqLength - sliceInfo_.sIdx;
689+ uint32_t tcGap;
690+ if (sliceInfo_.sIdx == 0) {
691+ tcGap = static_cast<uint32_t>(
692+ CeilDivT(sliceInfo_.bStartPos + seqLength, static_cast<uint64_t>(cmpRatio)) -
693+ sliceInfo_.bStartPos / static_cast<uint64_t>(cmpRatio));
694+ } else {
695+ tcGap = static_cast<uint32_t>(
696+ CeilDivT(sliceInfo_.bStartPos + seqLength, static_cast<uint64_t>(cmpRatio)) -
697+ CeilDivT(sliceInfo_.bStartPos + sliceInfo_.sIdx, static_cast<uint64_t>(cmpRatio)));
698+ }
699+ 
700+ if (needDealTcSize_ < tcGap) {
701+ // ── (B) 部分跳过：任务量不足以跳过整个空洞 ──
702+ // 只推进任务量对应的行（needTc 个块 = needTc*cmpRatio 行），clamp 到空洞内。
703+ //
704+ // ★特殊4：needTc == 0 时不推进任何行——行推进必须与任务量消耗严格
705+ // 对应；否则核起点"行已推进、块未消耗"，后续核分到无行的块
706+ // （读窗口外数据 / 输出丢失）。
707+ //
708+ // ★特殊5：needTc*cmpRatio 可能超过空洞行数（尾部块凑不齐一块），
709+ // clamp 后停在空洞末尾，不越界（"凑不齐也算一块"的任务量不变）。
710+ uint32_t skip = min(needDealTcSize_ * cmpRatio, gapRows);
711+ sliceInfo_.dealedSeqCnt += skip;
712+ sliceInfo_.sIdx += skip;
713+ needDealTcSize_ = 0;
714+ break; // 任务量耗尽：迭代终止
715+ }
716+ // ── (C) 完整跳过：推进整个空洞，消耗 tcGap 个 Tc ──
717+ // 空洞全部行在本 batch 内，推进后 sIdx 到达 seqLength（batch 末尾），
718+ // 随后换到下一个 batch。
719+ sliceInfo_.dealedSeqCnt += gapRows;
720+ sliceInfo_.sIdx += gapRows;
721+ needDealTcSize_ -= tcGap;
679 }722 }
680 sliceInfo_.bIdx++;723 sliceInfo_.bIdx++;
681 if (sliceInfo_.bIdx == batch_size_) {724 if (sliceInfo_.bIdx == batch_size_) {
682- sliceInfo_.bIdx = 0;725+ // 终止而非回绕（防死循环；正常遍历不会触发）
726+ sliceInfo_.bIdx = batch_size_ - 1;
727+ sliceInfo_.sIdx = 0;
728+ sliceInfo_.bSeqUsed = 0;
729+ sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
730+ sliceInfo_.bSeqLength = tools_.GetSeqLength(sliceInfo_.bIdx);
731+ needDealTcSize_ = 0;
732+ break;
683 }733 }
684 sliceInfo_.sIdx = 0;734 sliceInfo_.sIdx = 0;
685 sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);735 sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
736+ sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
737+ sliceInfo_.bSeqLength = tools_.GetSeqLength(sliceInfo_.bIdx);
686 } while (sliceInfo_.bSeqUsed == 0);738 } while (sliceInfo_.bSeqUsed == 0);
687- sliceInfo_.bSeqLength = tools_.GetSeqLength(sliceInfo_.bIdx);
688- sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
689 }739 }
690 if (isFirst_) {740 if (isFirst_) {
691 isFirst_ = false;741 isFirst_ = false;
@@ -709,7 +759,10 @@ template <typename COMP>
709__aicore__ inline Vec1SliceInfo &CompressorVec1SliceIterator<COMP>::GetSlice()759__aicore__ inline Vec1SliceInfo &CompressorVec1SliceIterator<COMP>::GetSlice()
710{760{
711 uint32_t cmpRatio = tools_.toolParams_.cmpRatio;761 uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
712- if (sliceInfo_.bSeqUsed < sliceInfo_.sIdx) {762+ if (sliceInfo_.bSeqUsed <= sliceInfo_.sIdx) {
763+ // sIdx == bSeqUsed 时同样视为无效 slice（含 bSeqUsed==0 的空洞 batch 起点）：
764+ // 否则 GetSlice 会生成 headHolder>0/valid=0 的伪 slice（dealTcSize=1）
765+ // 每轮白吃 1 个 tc 配额且不推进 → 迭代器卡死在空洞 batch，其后块全部丢失
713 sliceInfo_.headHolderSeqCnt = 0;766 sliceInfo_.headHolderSeqCnt = 0;
714 sliceInfo_.validSeqCnt = 0;767 sliceInfo_.validSeqCnt = 0;
715 sliceInfo_.tailHolderSeqCnt = 0;768 sliceInfo_.tailHolderSeqCnt = 0;
Mattention/compressor/op_kernel/arch35/compressor_block_vec.h+69-3文件内容审核中，请稍后刷新重试
Mattention/compressor/op_kernel/arch35/compressor_block_vec_full_load.h+69-7文件内容审核中，请稍后刷新重试
Mattention/compressor/op_kernel/arch35/compressor_comm.h+6-1
@@ -86,14 +86,19 @@ enum class CACHE_MODE : std::uint8_t {
86 RING_BUFFER = static_cast<std::uint8_t>(2)86 RING_BUFFER = static_cast<std::uint8_t>(2)
87};87};
88 88 
89+enum class GRAD_ENABLED : std::uint8_t {
90+ DISABLE = static_cast<std::uint8_t>(0),
91+ ENABLE = static_cast<std::uint8_t>(1)
92+};
89enum class TEMPLATE_ID : uint8_t { NORMAL = 0, EMPTY_X = 1, FULL_LOAD = 2 };93enum class TEMPLATE_ID : uint8_t { NORMAL = 0, EMPTY_X = 1, FULL_LOAD = 2 };
90 94 
91-template <X_LAYOUT X_L, X_DTYPE X_T, COFF C, CACHE_MODE Cache_Mode, typename... Args>95+template <X_LAYOUT X_L, X_DTYPE X_T, COFF C, CACHE_MODE Cache_Mode, GRAD_ENABLED G, typename... Args>
92struct COMPType {96struct COMPType {
93 static constexpr X_LAYOUT xLayout = X_L;97 static constexpr X_LAYOUT xLayout = X_L;
94 static constexpr X_DTYPE xDtype = X_T;98 static constexpr X_DTYPE xDtype = X_T;
95 static constexpr COFF coff = C;99 static constexpr COFF coff = C;
96 static constexpr CACHE_MODE cacheMode = Cache_Mode;100 static constexpr CACHE_MODE cacheMode = Cache_Mode;
101+ static constexpr GRAD_ENABLED gradEnabled = G;
97};102};
98 103 
99struct CmpBlockInfo {104struct CmpBlockInfo {
Mattention/compressor/op_kernel/arch35/compressor_kernel.h+6-3
@@ -37,7 +37,8 @@ public:
37 __aicore__ inline void Init(__gm__ uint8_t *x, __gm__ uint8_t *wKv, __gm__ uint8_t *wGate,37 __aicore__ inline void Init(__gm__ uint8_t *x, __gm__ uint8_t *wKv, __gm__ uint8_t *wGate,
38 __gm__ uint8_t *stateCache, __gm__ uint8_t *ape, __gm__ uint8_t *stateBlockTable,38 __gm__ uint8_t *stateCache, __gm__ uint8_t *ape, __gm__ uint8_t *stateBlockTable,
39 __gm__ uint8_t *cuSeqlens, __gm__ uint8_t *seqUsed, __gm__ uint8_t *startPos,39 __gm__ uint8_t *cuSeqlens, __gm__ uint8_t *seqUsed, __gm__ uint8_t *startPos,
40- __gm__ uint8_t *cmpKvOut, __gm__ uint8_t *workspace);40+ __gm__ uint8_t *cmpKvOut, __gm__ uint8_t *softmaxScoreOut, __gm__ uint8_t *kvOut,
41+ __gm__ uint8_t *workspace);
41 __aicore__ inline void Process();42 __aicore__ inline void Process();
42 43 
43private:44private:
@@ -110,7 +111,8 @@ __aicore__ inline void CompressorKernel<COMP>::Init(__gm__ uint8_t *x, __gm__ ui
110 __gm__ uint8_t *stateCache, __gm__ uint8_t *ape,111 __gm__ uint8_t *stateCache, __gm__ uint8_t *ape,
111 __gm__ uint8_t *stateBlockTable, __gm__ uint8_t *cuSeqlens,112 __gm__ uint8_t *stateBlockTable, __gm__ uint8_t *cuSeqlens,
112 __gm__ uint8_t *seqUsed, __gm__ uint8_t *startPos,113 __gm__ uint8_t *seqUsed, __gm__ uint8_t *startPos,
113- __gm__ uint8_t *cmpKvOut, __gm__ uint8_t *workspace)114+ __gm__ uint8_t *cmpKvOut, __gm__ uint8_t *softmaxScoreOut,
115+ __gm__ uint8_t *kvOut, __gm__ uint8_t *workspace)
114{116{
115 if ASCEND_IS_AIV {117 if ASCEND_IS_AIV {
116 constInfo.aiCoreIdx = GetBlockIdx() / 2;118 constInfo.aiCoreIdx = GetBlockIdx() / 2;
@@ -153,7 +155,8 @@ __aicore__ inline void CompressorKernel<COMP>::Init(__gm__ uint8_t *x, __gm__ ui
153 blockCube_.InitGlobalBuffers(mm1KvResGm, mm1ScoreResGm);155 blockCube_.InitGlobalBuffers(mm1KvResGm, mm1ScoreResGm);
154 } else {156 } else {
155 blockVec_.InitParams(constInfo, tools_);157 blockVec_.InitParams(constInfo, tools_);
156- blockVec_.Init(x, wKv, wGate, stateCache, ape, stateBlockTable, cuSeqlens, seqUsed, startPos, cmpKvOut);158+ blockVec_.Init(x, wKv, wGate, stateCache, ape, stateBlockTable, cuSeqlens,
159+ seqUsed, startPos, cmpKvOut, softmaxScoreOut, kvOut);
157 blockVec_.InitBuffers(pipe_);160 blockVec_.InitBuffers(pipe_);
158 blockVec_.InitVec1GlobalTensor(Vec1InputKvGm, Vec1InputScoreGm, vec1KvCacheGm, vec1ScoreCacheGm);161 blockVec_.InitVec1GlobalTensor(Vec1InputKvGm, Vec1InputScoreGm, vec1KvCacheGm, vec1ScoreCacheGm);
159 }162 }
Mattention/compressor/op_kernel/arch35/compressor_kernel_full_load.h+7-4
@@ -36,7 +36,8 @@ public:
36 __aicore__ inline void Init(__gm__ uint8_t *x, __gm__ uint8_t *wKv, __gm__ uint8_t *wGate,36 __aicore__ inline void Init(__gm__ uint8_t *x, __gm__ uint8_t *wKv, __gm__ uint8_t *wGate,
37 __gm__ uint8_t *stateCache, __gm__ uint8_t *ape, __gm__ uint8_t *stateBlockTable,37 __gm__ uint8_t *stateCache, __gm__ uint8_t *ape, __gm__ uint8_t *stateBlockTable,
38 __gm__ uint8_t *cuSeqlens, __gm__ uint8_t *seqUsed, __gm__ uint8_t *startPos,38 __gm__ uint8_t *cuSeqlens, __gm__ uint8_t *seqUsed, __gm__ uint8_t *startPos,
39- __gm__ uint8_t *cmpKvOut, __gm__ uint8_t *workspace);39+ __gm__ uint8_t *cmpKvOut, __gm__ uint8_t *softmaxScoreOut,
40+ __gm__ uint8_t *kvOut, __gm__ uint8_t *workspace);
40 __aicore__ inline void Process();41 __aicore__ inline void Process();
41 42 
42private:43private:
@@ -103,8 +104,9 @@ __aicore__ inline void CompressorKernelFullLoad<COMP>::Init(__gm__ uint8_t *x, _
103 __gm__ uint8_t *wGate, __gm__ uint8_t *stateCache,104 __gm__ uint8_t *wGate, __gm__ uint8_t *stateCache,
104 __gm__ uint8_t *ape, __gm__ uint8_t *stateBlockTable,105 __gm__ uint8_t *ape, __gm__ uint8_t *stateBlockTable,
105 __gm__ uint8_t *cuSeqlens, __gm__ uint8_t *seqUsed,106 __gm__ uint8_t *cuSeqlens, __gm__ uint8_t *seqUsed,
106- __gm__ uint8_t *startPos, __gm__ uint8_t *cmpKvOut,107+ __gm__ uint8_t *startPos,
107- __gm__ uint8_t *workspace)108+ __gm__ uint8_t *cmpKvOut, __gm__ uint8_t *softmaxScoreOut,
109+ __gm__ uint8_t *kvOut, __gm__ uint8_t *workspace)
108{110{
109 if ASCEND_IS_AIV {111 if ASCEND_IS_AIV {
110 constInfo.aiCoreIdx = GetBlockIdx() / 2;112 constInfo.aiCoreIdx = GetBlockIdx() / 2;
@@ -147,7 +149,8 @@ __aicore__ inline void CompressorKernelFullLoad<COMP>::Init(__gm__ uint8_t *x, _
147 blockCube_.InitGlobalBuffers(mm1KvResGm, mm1ScoreResGm);149 blockCube_.InitGlobalBuffers(mm1KvResGm, mm1ScoreResGm);
148 } else {150 } else {
149 blockVec_.InitParams(constInfo, tools_);151 blockVec_.InitParams(constInfo, tools_);
150- blockVec_.Init(x, wKv, wGate, stateCache, ape, stateBlockTable, cuSeqlens, seqUsed, startPos, cmpKvOut);152+ blockVec_.Init(x, wKv, wGate, stateCache, ape, stateBlockTable,
153+ cuSeqlens, seqUsed, startPos, cmpKvOut, softmaxScoreOut, kvOut);
151 blockVec_.InitBuffers(pipe_);154 blockVec_.InitBuffers(pipe_);
152 blockVec_.InitVec1GlobalTensor(Vec1InputKvGm, Vec1InputScoreGm, vec1KvCacheGm, vec1ScoreCacheGm);155 blockVec_.InitVec1GlobalTensor(Vec1InputKvGm, Vec1InputScoreGm, vec1KvCacheGm, vec1ScoreCacheGm);
153 }156 }
Mattention/compressor/op_kernel/arch35/compressor_template_tiling_key.h+3-0
@@ -35,6 +35,8 @@ ASCENDC_TPL_ARGS_DECL(compressor, // 算子唯一标识，与opType保持一致
35 ASCENDC_TPL_UINT_DECL(CACHE_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2),35 ASCENDC_TPL_UINT_DECL(CACHE_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2),
36 // bit:9-10 template_id 0:empty_tensor 1:normal 2:full load36 // bit:9-10 template_id 0:empty_tensor 1:normal 2:full load
37 ASCENDC_TPL_UINT_DECL(TEMPLATE_ID, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 0, 1, 2),37 ASCENDC_TPL_UINT_DECL(TEMPLATE_ID, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 0, 1, 2),
38+ // bit:11 grad_enabled 0:关闭 1:开启
39+ ASCENDC_TPL_BOOL_DECL(GRAD_ENABLED, 0, 1),
38 40 
39);41);
40 42 
@@ -45,6 +47,7 @@ ASCENDC_TPL_SEL(
45 ASCENDC_TPL_UINT_SEL(COFF, ASCENDC_TPL_UI_LIST, 1, 2),47 ASCENDC_TPL_UINT_SEL(COFF, ASCENDC_TPL_UI_LIST, 1, 2),
46 ASCENDC_TPL_UINT_SEL(CACHE_MODE, ASCENDC_TPL_UI_LIST, 1, 2),48 ASCENDC_TPL_UINT_SEL(CACHE_MODE, ASCENDC_TPL_UI_LIST, 1, 2),
47 ASCENDC_TPL_UINT_SEL(TEMPLATE_ID, ASCENDC_TPL_UI_LIST, 0, 1, 2),49 ASCENDC_TPL_UINT_SEL(TEMPLATE_ID, ASCENDC_TPL_UI_LIST, 0, 1, 2),
50+ ASCENDC_TPL_BOOL_SEL(GRAD_ENABLED, 0, 1),
48 ASCENDC_TPL_TILING_STRUCT_SEL(optiling::CompressorTilingData)), );51 ASCENDC_TPL_TILING_STRUCT_SEL(optiling::CompressorTilingData)), );
49 52 
50#endif // COMPRESSOR_TEMPLATE_TILING_KEY_H53#endif // COMPRESSOR_TEMPLATE_TILING_KEY_H
Mattention/compressor/op_kernel/arch35/compressor_tools.h+76-18文件内容审核中，请稍后刷新重试
Mattention/compressor/op_kernel/compressor.cpp+17-3
@@ -24,6 +24,7 @@
24 24
25using namespace Compressor;25using namespace Compressor;
26 26 
27+#if (__CCE_AICORE__ == 220)
27#define INVOKE_COMPRESSOR_GENERAL_OP_IMPL(templateClass, ...) \28#define INVOKE_COMPRESSOR_GENERAL_OP_IMPL(templateClass, ...) \
28 do { \29 do { \
29 templateClass<COMPType<__VA_ARGS__>> op(&pipe, tilingData); \30 templateClass<COMPType<__VA_ARGS__>> op(&pipe, tilingData); \
@@ -31,8 +32,18 @@ using namespace Compressor;
31 cuSeqlens, seqUsed, startPos, cmpKvOut, workspace); \32 cuSeqlens, seqUsed, startPos, cmpKvOut, workspace); \
32 op.Process(); \33 op.Process(); \
33 } while (0)34 } while (0)
35+#else
36+#define INVOKE_COMPRESSOR_GENERAL_OP_IMPL(templateClass, ...) \
37+ do { \
38+ templateClass<COMPType<__VA_ARGS__>> op(&pipe, tilingData); \
39+ op.Init(x, wKv, wGate, stateCache, ape, stateBlockTable, \
40+ cuSeqlens, seqUsed, startPos, cmpKvOut, \
41+ softmaxScoreOut, kvOut, workspace); \
42+ op.Process(); \
43+ } while (0)
44+#endif
34 45 
35-template<uint8_t XLayout, uint8_t XDType, uint8_t Coff, uint8_t CacheMode, uint8_t TemplateId>46+template<uint8_t XLayout, uint8_t XDType, uint8_t Coff, uint8_t CacheMode, uint8_t TemplateId, uint8_t GradEnabled>
36__global__ __aicore__ void compressor(47__global__ __aicore__ void compressor(
37 __gm__ uint8_t *x,48 __gm__ uint8_t *x,
38 __gm__ uint8_t *wKv,49 __gm__ uint8_t *wKv,
@@ -45,6 +56,8 @@ __global__ __aicore__ void compressor(
45 __gm__ uint8_t *startPos,56 __gm__ uint8_t *startPos,
46 __gm__ uint8_t *cmpKvOut,57 __gm__ uint8_t *cmpKvOut,
47 __gm__ uint8_t *stateCacheOut,58 __gm__ uint8_t *stateCacheOut,
59+ __gm__ uint8_t *softmaxScoreOut,
60+ __gm__ uint8_t *kvOut,
48 __gm__ uint8_t *workspace,61 __gm__ uint8_t *workspace,
49 __gm__ uint8_t *tiling) {62 __gm__ uint8_t *tiling) {
50 REGISTER_TILING_DEFAULT(optiling::CompressorTilingData);63 REGISTER_TILING_DEFAULT(optiling::CompressorTilingData);
@@ -66,10 +79,11 @@ __global__ __aicore__ void compressor(
66 INVOKE_COMPRESSOR_GENERAL_OP_IMPL(CompressorKernelPerf, xLayout, xDtype, coff);79 INVOKE_COMPRESSOR_GENERAL_OP_IMPL(CompressorKernelPerf, xLayout, xDtype, coff);
67 }80 }
68 #else81 #else
82+ constexpr auto gradEnabled = static_cast<GRAD_ENABLED>(GradEnabled);
69 if constexpr (static_cast<TEMPLATE_ID>(TemplateId) == TEMPLATE_ID::FULL_LOAD) {83 if constexpr (static_cast<TEMPLATE_ID>(TemplateId) == TEMPLATE_ID::FULL_LOAD) {
70- INVOKE_COMPRESSOR_GENERAL_OP_IMPL(CompressorKernelFullLoad, xLayout, xDtype, coff, cacheMode);84+ INVOKE_COMPRESSOR_GENERAL_OP_IMPL(CompressorKernelFullLoad, xLayout, xDtype, coff, cacheMode, gradEnabled);
71 } else {85 } else {
72- INVOKE_COMPRESSOR_GENERAL_OP_IMPL(CompressorKernel, xLayout, xDtype, coff, cacheMode);86+ INVOKE_COMPRESSOR_GENERAL_OP_IMPL(CompressorKernel, xLayout, xDtype, coff, cacheMode, gradEnabled);
73 }87 }
74 #endif88 #endif
75}89}
Mattention/compressor/tests/ut/op_host/arch35/test_compressor_tiling.cpp+20-0
@@ -76,6 +76,8 @@ TEST_F(CompressorTilingArch35, test1)
76 {76 {
77 {{{2, 2, 512}, {2, 2, 512}}, ge::DT_BF16, ge::FORMAT_ND}, // cmp_kv77 {{{2, 2, 512}, {2, 2, 512}}, ge::DT_BF16, ge::FORMAT_ND}, // cmp_kv
78 {{{4, 128, 2048}, {4, 128, 2048}}, ge::DT_FLOAT, ge::FORMAT_ND}, // state_cache (in-place)78 {{{4, 128, 2048}, {4, 128, 2048}}, ge::DT_FLOAT, ge::FORMAT_ND}, // state_cache (in-place)
79+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
80+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
79 },81 },
80 {82 {
81 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},83 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
@@ -515,6 +517,8 @@ TEST_F(CompressorTilingArch35, test2)
515 {517 {
516 {{{2, 2, 512}, {2, 2, 512}}, ge::DT_FLOAT16, ge::FORMAT_ND},518 {{{2, 2, 512}, {2, 2, 512}}, ge::DT_FLOAT16, ge::FORMAT_ND},
517 {{{4, 128, 2048}, {4, 128, 2048}}, ge::DT_FLOAT, ge::FORMAT_ND},519 {{{4, 128, 2048}, {4, 128, 2048}}, ge::DT_FLOAT, ge::FORMAT_ND},
520+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
521+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
518 },522 },
519 {523 {
520 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},524 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
@@ -547,6 +551,8 @@ TEST_F(CompressorTilingArch35, test3)
547 {551 {
548 {{{2, 2, 128}, {2, 2, 128}}, ge::DT_BF16, ge::FORMAT_ND},552 {{{2, 2, 128}, {2, 2, 128}}, ge::DT_BF16, ge::FORMAT_ND},
549 {{{4, 128, 512}, {4, 128, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},553 {{{4, 128, 512}, {4, 128, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
554+ {{{2, 2, 8, 128}, {2, 2, 8, 128}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
555+ {{{2, 2, 8, 128}, {2, 2, 8, 128}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
550 },556 },
551 {557 {
552 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},558 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
@@ -579,6 +585,8 @@ TEST_F(CompressorTilingArch35, test4)
579 {585 {
580 {{{2, 1, 512}, {2, 1, 512}}, ge::DT_BF16, ge::FORMAT_ND},586 {{{2, 1, 512}, {2, 1, 512}}, ge::DT_BF16, ge::FORMAT_ND},
581 {{{4, 128, 1024}, {4, 128, 1024}}, ge::DT_FLOAT, ge::FORMAT_ND},587 {{{4, 128, 1024}, {4, 128, 1024}}, ge::DT_FLOAT, ge::FORMAT_ND},
588+ {{{2, 1, 128, 512}, {2, 1, 128, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
589+ {{{2, 1, 128, 512}, {2, 1, 128, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
582 },590 },
583 {591 {
584 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(128)},592 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(128)},
@@ -611,6 +619,8 @@ TEST_F(CompressorTilingArch35, test5)
611 {619 {
612 {{{1, 2, 128}, {1, 2, 128}}, ge::DT_BF16, ge::FORMAT_ND},620 {{{1, 2, 128}, {1, 2, 128}}, ge::DT_BF16, ge::FORMAT_ND},
613 {{{2, 128, 512}, {2, 128, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},621 {{{2, 128, 512}, {2, 128, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
622+ {{{1, 2, 8, 128}, {1, 2, 8, 128}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
623+ {{{1, 2, 8, 128}, {1, 2, 8, 128}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
614 },624 },
615 {625 {
616 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},626 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
@@ -647,6 +657,8 @@ TEST_F(CompressorTilingArch35, test6)
647 {657 {
648 {{{2, 512}, {2, 512}}, ge::DT_BF16, ge::FORMAT_ND},658 {{{2, 512}, {2, 512}}, ge::DT_BF16, ge::FORMAT_ND},
649 {{{4, 128, 2048}, {4, 128, 2048}}, ge::DT_FLOAT, ge::FORMAT_ND},659 {{{4, 128, 2048}, {4, 128, 2048}}, ge::DT_FLOAT, ge::FORMAT_ND},
660+ {{{2, 8, 512}, {2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
661+ {{{2, 8, 512}, {2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
650 },662 },
651 {663 {
652 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},664 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
@@ -683,6 +695,8 @@ TEST_F(CompressorTilingArch35, test7)
683 {695 {
684 {{{2, 2, 512}, {2, 2, 512}}, ge::DT_BF16, ge::FORMAT_ND},696 {{{2, 2, 512}, {2, 2, 512}}, ge::DT_BF16, ge::FORMAT_ND},
685 {{{4, 128, 2048}, {4, 128, 2048}}, ge::DT_FLOAT, ge::FORMAT_ND},697 {{{4, 128, 2048}, {4, 128, 2048}}, ge::DT_FLOAT, ge::FORMAT_ND},
698+ {{{2, 3, 8, 512}, {2, 3, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
699+ {{{2, 3, 8, 512}, {2, 3, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
686 },700 },
687 {701 {
688 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(3)},702 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(3)},
@@ -715,6 +729,8 @@ TEST_F(CompressorTilingArch35, test8)
715 {729 {
716 {{{2, 2, 512}, {2, 2, 512}}, ge::DT_BF16, ge::FORMAT_ND},730 {{{2, 2, 512}, {2, 2, 512}}, ge::DT_BF16, ge::FORMAT_ND},
717 {{{4, 128, 2048}, {4, 128, 2048}}, ge::DT_FLOAT, ge::FORMAT_ND},731 {{{4, 128, 2048}, {4, 128, 2048}}, ge::DT_FLOAT, ge::FORMAT_ND},
732+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
733+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
718 },734 },
719 {735 {
720 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},736 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
@@ -747,6 +763,8 @@ TEST_F(CompressorTilingArch35, test9)
747 {763 {
748 {{{2, 2, 256}, {2, 2, 256}}, ge::DT_BF16, ge::FORMAT_ND},764 {{{2, 2, 256}, {2, 2, 256}}, ge::DT_BF16, ge::FORMAT_ND},
749 {{{4, 128, 1024}, {4, 128, 1024}}, ge::DT_FLOAT, ge::FORMAT_ND},765 {{{4, 128, 1024}, {4, 128, 1024}}, ge::DT_FLOAT, ge::FORMAT_ND},
766+ {{{2, 2, 8, 256}, {2, 2, 8, 256}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
767+ {{{2, 2, 8, 256}, {2, 2, 8, 256}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
750 },768 },
751 {769 {
752 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},770 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
@@ -779,6 +797,8 @@ TEST_F(CompressorTilingArch35, test10)
779 {797 {
780 {{{2, 2, 512}, {2, 2, 512}}, ge::DT_BF16, ge::FORMAT_ND},798 {{{2, 2, 512}, {2, 2, 512}}, ge::DT_BF16, ge::FORMAT_ND},
781 {{{4, 128, 2048}, {4, 128, 2048}}, ge::DT_FLOAT, ge::FORMAT_ND},799 {{{4, 128, 2048}, {4, 128, 2048}}, ge::DT_FLOAT, ge::FORMAT_ND},
800+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
801+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
782 },802 },
783 {803 {
784 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},804 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
Mattention/compressor/tests/ut/op_host/test_compressor_infershape.cpp+40-16
@@ -56,14 +56,17 @@ TEST_F(CompressorInfershape, bsh_c4a_bf16)
56 },56 },
57 {57 {
58 // output 0: cmp_kv (placeholder)58 // output 0: cmp_kv (placeholder)
59- {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND},59+ {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND}, {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // state_cache
60+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
61+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
62+ 
60 },63 },
61 {64 {
62 // attrs65 // attrs
63 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},66 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
64 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},67 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
65 });68 });
66- std::vector<std::vector<int64_t>> expectOutputShape = {{2, 1, 512}};69+ std::vector<std::vector<int64_t>> expectOutputShape = {{2, 1, 512}, {}, {2, 1, 8, 512}, {2, 1, 8, 512}};
67 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);70 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);
68}71}
69 72 
@@ -83,13 +86,16 @@ TEST_F(CompressorInfershape, bsh_c4a_fp16)
83 {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},86 {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
84 },87 },
85 {88 {
86- {{{}, {}}, ge::DT_FLOAT16, ge::FORMAT_ND},89+ {{{}, {}}, ge::DT_FLOAT16, ge::FORMAT_ND}, {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // state_cache
90+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
91+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
92+ 
87 },93 },
88 {94 {
89 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},95 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
90 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},96 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
91 });97 });
92- std::vector<std::vector<int64_t>> expectOutputShape = {{2, 1, 512}};98+ std::vector<std::vector<int64_t>> expectOutputShape = {{2, 1, 512}, {}, {2, 1, 8, 512}, {2, 1, 8, 512}};
93 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);99 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);
94}100}
95 101 
@@ -110,13 +116,16 @@ TEST_F(CompressorInfershape, bsh_c4li_bf16)
110 {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},116 {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
111 },117 },
112 {118 {
113- {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND},119+ {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND}, {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // state_cache
120+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
121+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
122+ 
114 },123 },
115 {124 {
116 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},125 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
117 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},126 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
118 });127 });
119- std::vector<std::vector<int64_t>> expectOutputShape = {{2, 1, 128}};128+ std::vector<std::vector<int64_t>> expectOutputShape = {{2, 1, 128}, {}, {2, 1, 8, 128}, {2, 1, 8, 128}};
120 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);129 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);
121}130}
122 131 
@@ -137,13 +146,16 @@ TEST_F(CompressorInfershape, bsh_c128a_bf16)
137 {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},146 {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
138 },147 },
139 {148 {
140- {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND},149+ {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND}, {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // state_cache
150+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
151+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
152+ 
141 },153 },
142 {154 {
143 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(128)},155 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(128)},
144 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},156 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},
145 });157 });
146- std::vector<std::vector<int64_t>> expectOutputShape = {{2, 1, 512}};158+ std::vector<std::vector<int64_t>> expectOutputShape = {{2, 1, 512}, {}, {2, 1, 128, 512}, {2, 1, 128, 512}};
147 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);159 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);
148}160}
149 161 
@@ -164,13 +176,16 @@ TEST_F(CompressorInfershape, bsh_non_divisible)
164 {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},176 {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
165 },177 },
166 {178 {
167- {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND},179+ {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND}, {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // state_cache
180+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
181+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
182+ 
168 },183 },
169 {184 {
170 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},185 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
171 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},186 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
172 });187 });
173- std::vector<std::vector<int64_t>> expectOutputShape = {{1, 2, 128}};188+ std::vector<std::vector<int64_t>> expectOutputShape = {{1, 2, 128}, {}, {1, 2, 8, 128}, {1, 2, 8, 128}};
174 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);189 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);
175}190}
176 191 
@@ -202,13 +217,16 @@ TEST_F(CompressorInfershape, th_c4a_bf16)
202 {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND}, // start_pos217 {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND}, // start_pos
203 },218 },
204 {219 {
205- {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND},220+ {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND}, {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // state_cache
221+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
222+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
223+ 
206 },224 },
207 {225 {
208 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},226 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
209 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},227 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
210 });228 });
211- std::vector<std::vector<int64_t>> expectOutputShape = {{4, 512}};229+ std::vector<std::vector<int64_t>> expectOutputShape = {{4, 512}, {}, {4, 8, 512}, {4, 8, 512}};
212 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);230 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);
213}231}
214 232 
@@ -229,13 +247,16 @@ TEST_F(CompressorInfershape, th_c4a_multi_batch)
229 {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},247 {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
230 },248 },
231 {249 {
232- {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND},250+ {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND}, {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // state_cache
251+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
252+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
253+ 
233 },254 },
234 {255 {
235 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},256 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
236 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},257 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
237 });258 });
238- std::vector<std::vector<int64_t>> expectOutputShape = {{7, 512}};259+ std::vector<std::vector<int64_t>> expectOutputShape = {{7, 512}, {}, {7, 8, 512}, {7, 8, 512}};
239 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);260 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);
240}261}
241 262 
@@ -258,14 +279,17 @@ TEST_F(CompressorInfershape, bsh_empty_batch)
258 {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},279 {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
259 },280 },
260 {281 {
261- {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND},282+ {{{}, {}}, ge::DT_BF16, ge::FORMAT_ND}, {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // state_cache
283+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
284+ {{{}, {}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
285+ 
262 },286 },
263 {287 {
264 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},288 {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
265 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},289 {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
266 });290 });
267 // B=0, Sr=ceil(4/4)=1, output: [0, 1, 512]291 // B=0, Sr=ceil(4/4)=1, output: [0, 1, 512]
268- std::vector<std::vector<int64_t>> expectOutputShape = {{0, 1, 512}};292+ std::vector<std::vector<int64_t>> expectOutputShape = {{0, 1, 512}, {}, {0, 1, 8, 512}, {0, 1, 8, 512}};
269 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);293 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);
270}294}
271 295 
Aattention/compressor_grad/CMakeLists.txt+18-0
@@ -0,0 +1,18 @@
1+# -----------------------------------------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# -----------------------------------------------------------------------------------------------------------
10+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
11+if(NOT ENABLE_TEST AND NOT BENCHMARK)
12+ list(REMOVE_ITEM CURRENT_DIRS tests)
13+endif()
14+foreach(SUB_DIR ${CURRENT_DIRS})
15+ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
16+ add_subdirectory(${SUB_DIR})
17+ endif()
18+endforeach()
Aattention/compressor_grad/README.md+109-0
@@ -0,0 +1,109 @@
1+# CompressorGrad
2+ 
3+## 产品支持情况
4+ 
5+| 产品 | 是否支持 |
6+| ------------------------------------------------------------ | :------: |
7+|<term>Ascend 950PR/Ascend 950DT</term>| √ |
8+|<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>| × |
9+|<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>| × |
10+|<term>Atlas 200I/500 A2 推理产品</term>| × |
11+|<term>Atlas 推理系列产品</term>| × |
12+|<term>Atlas 训练系列产品</term>| × |
13+ 
14+## 功能说明
15+ 
16+- 算子功能：CompressorGrad是Compressor算子的反向算子，用于计算输入$X$、权重$W^{KV}$/$W^{Gate}$与位置编码$Ape$的梯度。前向在gradEnabled为true时导出softmax\_score（分组softmax结果）与kv（softmax结果与kv\_state的Hadamard乘积）中间结果，作为本算子的输入。主要计算过程为：
17+ 1. 逐块计算Hadamard积反向：将上游梯度$dC$与softmax\_score、kv逐元素相乘，得到$dK$与$dS^\prime$；
18+ 2. softmax反向：对$dS^\prime$沿压缩轴做softmax反向，得到$dZ$；
19+ 3. APE梯度计算：按token位置累加$dZ$，得到$dApe$；
20+ 4. matmul反向：将$dK$、$dZ$与权重做矩阵乘法反向，得到$dX$、$dW^{KV}$、$dW^{Gate}$。
21+ 
22+- 计算公式：
23+ 
24+ 1. 计算Hadamard乘积反向，$N$为压缩块总数，$i$为压缩块序号，$dC_i$为第$i$块的上游梯度，$S_i$、$K_i$分别为softmax\_score、kv第$i$块：
25+ 
26+ $$
27+ dK_i = dC_i \odot S_i,~ i=1,\cdots,N
28+ $$
29+ 
30+ $$
31+ dS^\prime_i = dC_i \odot K_i,~ i=1,\cdots,N
32+ $$
33+ 
34+ 2. 计算softmax反向（沿压缩轴求和），$k$为块内行序号：
35+ 
36+ $$
37+ dZ_i = S_i \odot \left(dS^\prime_i - \sum_{k=1}^{coff \cdot cmp\_ratio} \left(S_i \odot dS^\prime_i\right)_{k,:}\right),~ i=1,\cdots,N
38+ $$
39+ 
40+ 3. 计算APE梯度，$pos$为$dZ$各行对应token的全局位置：
41+ 
42+ $$
43+ dApe = ScatterAdd\left(dZ,~ pos \% cmp\_ratio\right)
44+ $$
45+ 
46+ 4. 计算矩阵乘法反向，$dNewKv$、$dNewScore$为$dK$、$dZ$按压缩块映射回全局token行的结果（coff=2时prev/cur半区分别对应上一块与本块的token行，与正向的$W^{aKV}$/$W^{bKV}$对应）：
47+ 
48+ $$
49+ dX = dNewKv @ W^{KV} + dNewScore @ W^{Gate}
50+ $$
51+ 
52+ $$
53+ dW^{KV} = dNewKv^T @ X,~ dW^{Gate} = dNewScore^T @ X
54+ $$
55+## 参数说明
56+ 
57+| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 |
58+|----------------------------|-----------|----------------------------------------------------------------------|----------------|------------|
59+| x | 输入 | 公式中的$X$，前向输入的原始数据。 | FLOAT16、BFLOAT16 | ND |
60+| wkv | 输入 | 公式中的$W^{KV}$，前向kv压缩权重。 | FLOAT16、BFLOAT16 | ND |
61+| wgate | 输入 | 公式中的$W^{Gate}$，前向gate压缩权重。 | FLOAT16、BFLOAT16 | ND |
62+| d\_cmp\_kv | 输入 | 公式中的$dC$，前向输出cmp\_kv的上游梯度。 | FLOAT16、BFLOAT16 | ND |
63+| softmax\_score | 输入 | 公式中的$S$，前向在gradEnabled为true时导出的分组softmax中间结果。 | FLOAT32 | ND |
64+| kv | 输入 | 公式中的$K$，前向在gradEnabled为true时导出的softmax结果与kv\_state的Hadamard乘积中间结果。 | FLOAT32 | ND |
65+| cu\_seqlens | 可选输入 | 表示不同Batch中的有效token数。<br>当x的shape为[T,H]时必传；当x的shape为[B,S,H]时，参数必须为空。 | INT32 | ND |
66+| seqused | 可选输入 | 表示不同Batch中实际参与压缩的token数。<br>如果指定为None时，表示和每个Batch上的Sequence Length长度相同。 | INT32 | ND |
67+| start\_pos | 可选输入 | 表示计算起始位置。<br>如果指定为None时，表示从0开始进行计算。 | INT32 | ND |
68+| cmp\_ratio | 属性 | 用于稀疏计算，表示数据压缩率，取值范围为[2, 128]内的整数。 | INT32 | - |
69+| coff | 可选属性 | 表示是否进行overlap数据重排。 <br>coff=1：无需进行overlap数据重排，coff=2：需要进行overlap数据重排。<br>默认值为1。 | INT32 | - |
70+| d\_x | 输出 | 公式中的$dX$，输入x的梯度。 | FLOAT16、BFLOAT16 | ND |
71+| d\_wkv | 输出 | 公式中的$dW^{KV}$，权重wkv的梯度。 | FLOAT16、BFLOAT16 | ND |
72+| d\_wgate | 输出 | 公式中的$dW^{Gate}$，权重wgate的梯度。 | FLOAT16、BFLOAT16 | ND |
73+| d\_ape | 输出 | 公式中的$dApe$，APE位置编码的梯度。 | FLOAT32 | ND |
74+ 
75+## 约束说明
76+ 
77+- x参数维度含义：B（Batch Size）表示输入样本批量大小、S（Sequence Length）表示输入样本序列长度、H（Head Size）表示hidden层的大小、D（Head Dim）表示hidden层的最小单元大小、T表示所有Batch输入样本序列长度的累加和。
78+- 输入shape限制：
79+ - wkv支持输入shape[coff* D,H]
80+ - wgate支持输入shape[coff* D,H]
81+ - softmax\_score支持输入shape：BS合轴时为[min(T,T//cmp_ratio+B), coff*cmp_ratio, D]；BS非合轴时为[B,ceil(S/cmp_ratio),coff*cmp_ratio,D]
82+ - kv支持输入shape：同softmax\_score
83+ - d\_cmp\_kv支持输入shape：BS合轴时为[min(T,T//cmp_ratio+B),D]；BS非合轴时为[B,ceil(S/cmp_ratio),D]
84+ - d\_x支持输出shape：与x相同，BS合轴时为[T,H]、BS非合轴时为[B,S,H]
85+ - d\_wkv、d\_wgate支持输出shape[coff* D,H]
86+ - d\_ape支持输出shape[cmp_ratio,coff* D]
87+ - start\_pos支持输入shape[B,]
88+ - 若x的维度采用BS合轴，即x的输入shape为[T,H]
89+ - cu\_seqlens输入shape必须为[B+1,]。该参数中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值，且第一位必须为0。
90+ - seqused，支持输入shape[B,]，要求每个Batch的有效token数要求小于等于对应Sequence Length长度，即seqused[n] <= cu\_seqlens[n+1] - cu\_seqlens[n]，且不小于0。
91+ - 若x的维度不采用BS合轴，即x的输入shape为[B,S,H]
92+ - cu\_seqlens，参数必须为空。
93+ - seqused，支持输入shape[B,]，要求每个Batch的有效token数要求小于等于对应Sequence Length长度，即要求seqused[n] <= S，且不小于0。
94+- 输入值域限制：
95+ - 该接口支持B、S泛化，且存在如下场景限制：
96+ - **不支持B、S、T为0的空Tensor**：与正向Compressor不同，CompressorGrad所有输入/输出均不支持空Tensor，shapeSize必须大于0。
97+ - 部分长序列场景下，如果计算量过大可能会导致出现超过NPU内存的报错，注：这里计算量会受x输入shape的影响，值越大计算量越大。
98+- 输入属性限制：
99+ - 支持D为128/512。
100+ - 支持H为1K~10K，512对齐。
101+ - 支持coff为1/2。
102+ - 支持cmp\_ratio为2~128。
103+ 
104+## 调用说明
105+ 
106+| 调用方式 | 样例代码 | 说明 |
107+| --------- | --------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
108+| aclnn接口 | -| 通过[aclnnCompressorGrad](./docs/aclnnCompressorGrad.md)调用CompressorGrad算子。 |
109+| PyTorch API | - | 通过[cann_ops_transformer.compressor](../../torch_extension/cann_ops_transformer/docs/zh/compressor.md)调用Compressor算子，反向经其autograd自动调用CompressorGrad。 |
Aattention/compressor_grad/docs/aclnnCompressorGrad.md+236-0
@@ -0,0 +1,236 @@
1+# aclnnCompressorGrad
C
Cchenjiao15 天前

交付件不全

likedislike
2+ 
3+## 产品支持情况
4+ 
5+<!-- npu="950" id1 -->
6+- <term>Ascend 950PR/Ascend 950DT</term>：支持
7+<!-- end id1 -->
8+<!-- npu="A3" id2 -->
9+- <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：不支持
10+<!-- end id2 -->
11+<!-- npu="910b" id3 -->
12+- <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>：不支持
13+<!-- end id3 -->
14+<!-- npu="310b" id4 -->
15+- <term>Atlas 200I/500 A2 推理产品</term>：不支持
16+<!-- end id4 -->
17+<!-- npu="310p" id5 -->
18+- <term>Atlas 推理系列产品</term>：不支持
19+<!-- end id5 -->
20+<!-- npu="910" id6 -->
21+- <term>Atlas 训练系列产品</term>：不支持
22+<!-- end id6 -->
23+ 
24+## 功能说明
25+ 
26+- 接口功能：CompressorGrad是Compressor算子的反向算子，用于计算输入$X$、权重$W^{KV}$/$W^{Gate}$与位置编码$Ape$的梯度。前向在gradEnabled为true时导出softmax\_score（分组softmax结果）与kv（softmax结果与kv\_state的Hadamard乘积）中间结果，作为本算子的输入。主要计算过程为：
27+ 1. 逐块计算Hadamard积反向：将上游梯度$dC$与softmax\_score、kv逐元素相乘，得到$dK$与$dS^\prime$；
28+ 2. softmax反向：对$dS^\prime$沿压缩轴做softmax反向，得到$dZ$；
29+ 3. APE梯度计算：按token位置累加$dZ$，得到$dApe$；
30+ 4. matmul反向：将$dK$、$dZ$与权重做矩阵乘法反向，得到$dX$、$dW^{KV}$、$dW^{Gate}$。
31+ 
32+- 计算公式：
33+ 
34+ 1. 计算Hadamard乘积反向，$N$为压缩块总数，$i$为压缩块序号，$dC_i$为第$i$块的上游梯度，$S_i$、$K_i$分别为softmax\_score、kv第$i$块：
35+ 
36+ $$
37+ dK_i = dC_i \odot S_i,~ i=1,\cdots,N
38+ $$
39+ 
40+ $$
41+ dS^\prime_i = dC_i \odot K_i,~ i=1,\cdots,N
42+ $$
43+ 
44+ 2. 计算softmax反向（沿压缩轴求和），$k$为块内行序号：
45+ 
46+ $$
47+ dZ_i = S_i \odot \left(dS^\prime_i - \sum_{k=1}^{coff \cdot cmp\_ratio} \left(S_i \odot dS^\prime_i\right)_{k,:}\right),~ i=1,\cdots,N
48+ $$
49+ 
50+ 3. 计算APE梯度，$pos$为$dZ$各行对应token的全局位置：
51+ 
52+ $$
53+ dApe = ScatterAdd\left(dZ,~ pos \% cmp\_ratio\right)
54+ $$
55+ 
56+ 4. 计算矩阵乘法反向，$dNewKv$、$dNewScore$为$dK$、$dZ$按压缩块映射回全局token行的结果（coff=2时prev/cur半区分别对应上一块与本块的token行，与正向的$W^{aKV}$/$W^{bKV}$对应）：
57+ 
58+ $$
59+ dX = dNewKv @ W^{KV} + dNewScore @ W^{Gate}
60+ $$
61+ 
62+ $$
63+ dW^{KV} = dNewKv^T @ X,~ dW^{Gate} = dNewScore^T @ X
64+ $$
65+## 函数原型
66+ 
67+每个算子分为[两段式接口](../../../docs/zh/context/two_phase_api.md)，必须先调用“aclnnCompressorGradGetWorkspaceSize”接口获取入参并根据流程计算所需workspace大小，再调用“aclnnCompressorGrad”接口执行计算。
68+ 
69+```cpp
70+aclnnStatus aclnnCompressorGradGetWorkspaceSize(
71+ const aclTensor *x,
72+ const aclTensor *wkv,
73+ const aclTensor *wgate,
74+ const aclTensor *dCmpKv,
75+ const aclTensor *softmaxScore,
76+ const aclTensor *kv,
77+ const aclTensor *cuSeqlensOptional,
78+ const aclTensor *sequsedOptional,
79+ const aclTensor *startPosOptional,
80+ int64_t cmpRatio,
81+ int64_t coff,
82+ const aclTensor *dXOut,
83+ const aclTensor *dWkvOut,
84+ const aclTensor *dWgateOut,
85+ const aclTensor *dApeOut,
86+ uint64_t *workspaceSize,
87+ aclOpExecutor **executor)
88+```
89+ 
90+``` cpp
91+aclnnStatus aclnnCompressorGrad(
92+ void *workspace,
93+ uint64_t workspaceSize,
94+ aclOpExecutor *executor,
95+ aclrtStream stream)
96+```
97+ 
98+## aclnnCompressorGradGetWorkspaceSize
99+ 
100+- **参数说明**
101+ 
102+ | 参数名 | 输入/输出 | 描述 | 使用说明 | 数据类型 | 数据格式 | 维度（shape） | 非连续Tensor |
103+ |----------------------------|-----------|----------------------------------------------------------------------|----------------|------------|-|-|-|
104+ | x | 输入 | 公式中的$X$，前向输入的原始数据。 | 不支持空Tensor。 | FLOAT16、BFLOAT16 | ND | BS合轴：[T,H]、BS非合轴：[B,S,H]|×|
105+ | wkv | 输入 | 公式中的$W^{KV}$，前向kv压缩权重。 |不支持空Tensor。| FLOAT16、BFLOAT16 | ND |[coff* D,H]|×|
106+ | wgate | 输入 | 公式中的$W^{Gate}$，前向gate压缩权重。 |不支持空Tensor。| FLOAT16、BFLOAT16 | ND |[coff* D,H]|×|
107+ | dCmpKv | 输入 | 公式中的$dC$，前向输出cmp\_kv的上游梯度。 | 不支持空Tensor。 | FLOAT16、BFLOAT16 | ND |BS合轴：[min(T,T//cmp_ratio+B),D]、BS非合轴：[B,ceil(S/cmp_ratio),D]|×|
108+ | softmaxScore | 输入 | 公式中的$S$，前向在gradEnabled为true时导出的分组softmax中间结果。 | 不支持空Tensor。 | FLOAT32 | ND |BS合轴：[min(T,T//cmp_ratio+B), coff*cmp_ratio, D]、BS非合轴：[B,ceil(S/cmp_ratio),coff*cmp_ratio,D]|×|
109+ | kv | 输入 | 公式中的$K$，前向在gradEnabled为true时导出的softmax结果与kv\_state的Hadamard乘积中间结果。 | 不支持空Tensor。 | FLOAT32 | ND |同softmaxScore|×|
110+ | cuSeqlensOptional | 可选输入 | 表示不同Batch中的有效token数。 | 当x的shape为[T,H]时必传，输入shape为[B+1,]；当x的shape为[B,S,H]时，参数必须为空。不支持空Tensor。| INT32 | ND |当x的shape为[T,H]时，输入shape为[B+1,]|×|
111+ | sequsedOptional | 可选输入 | 表示不同Batch中实际参与压缩的token数。 | 为None时，表示和每个Batch上的Sequence Length长度相同；要求seqused[n]不超过对应Sequence Length，且不小于0。| INT32 | ND |[B,]|×|
112+ | startPosOptional | 可选输入 | 表示计算起始位置。 | 为None时，表示从0开始进行计算。| INT32 | ND |[B,]|×|
113+ | cmpRatio | 输入 | 用于稀疏计算，表示数据压缩率，与前向一致。 |取值范围为[2, 128]内的整数。| INT32 | - |-|-|
114+ | coff | 可选输入 | 表示是否进行overlap数据重排，与前向一致。 |取值范围为[1, 2]。当coff=1时，无需进行overlap数据重排。当coff=2时，需要进行overlap数据重排。| INT32 | - |-|-|
115+ | dXOut | 输出 | 公式中的$dX$，输入x的梯度。 | 不支持空Tensor。| FLOAT16、BFLOAT16 | ND |与x相同：BS合轴：[T,H]、BS非合轴：[B,S,H]|×|
116+ | dWkvOut | 输出 | 公式中的$dW^{KV}$，权重wkv的梯度。 | 不支持空Tensor。| FLOAT16、BFLOAT16 | ND |[coff* D,H]|×|
117+ | dWgateOut | 输出 | 公式中的$dW^{Gate}$，权重wgate的梯度。 | 不支持空Tensor。| FLOAT16、BFLOAT16 | ND |[coff* D,H]|×|
118+ | dApeOut | 输出 | 公式中的$dApe$，APE位置编码的梯度。 | 不支持空Tensor。| FLOAT32 | ND |[cmp_ratio,coff* D]|×|
119+ 
120+- **返回值**
121+ 
122+ aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn_return_code.md)。
123+ 
124+ 第一段接口完成入参校验，出现以下场景时报错：
125+ 
126+ <table style="undefined;table-layout: fixed;width: 1155px"><colgroup>
127+ <col style="width: 319px">
128+ <col style="width: 144px">
129+ <col style="width: 671px">
130+ </colgroup>
131+ <thead>
132+ <tr>
133+ <th>返回值</th>
134+ <th>错误码</th>
135+ <th>描述</th>
136+ </tr>
137+ </thead>
138+ <tbody>
139+ <tr>
140+ <td>ACLNN_ERR_PARAM_NULLPTR</td>
141+ <td>161001</td>
142+ <td>必须传入的参数（如接口核心依赖的输入/输出参数）中存在空指针。</td>
143+ </tr>
144+ <tr>
145+ <td>ACLNN_ERR_PARAM_INVALID</td>
146+ <td>161002</td>
147+ <td>输入参数的shape（维度/尺寸）、dtype（数据类型）不在接口支持的范围内。</td>
148+ </tr>
149+ <tr>
150+ <td>ACLNN_ERR_RUNTIME_ERROR</td>
151+ <td>361001</td>
152+ <td>API内存调用NPU Runtime接口时发生异常（如Runtime服务未启动、内存申请失败等）。</td>
153+ </tr>
154+ <tr>
155+ <td>ACLNN_ERR_INNER_TILING_ERROR</td>
156+ <td>561002</td>
157+ <td>tiling发生异常，入参的dtype类型或者shape错误。</td>
158+ </tr>
159+ </tbody>
160+ </table>
161+ 
162+## aclnnCompressorGrad
163+ 
164+- **参数说明**
165+ 
166+ <table style="undefined;table-layout: fixed; width: 1154px"><colgroup>
167+ <col style="width: 153px">
168+ <col style="width: 121px">
169+ <col style="width: 880px">
170+ </colgroup>
171+ <thead>
172+ <tr>
173+ <th>参数名</th>
174+ <th>输入/输出</th>
175+ <th>描述</th>
176+ </tr></thead>
177+ <tbody>
178+ <tr>
179+ <td>workspace</td>
180+ <td>输入</td>
181+ <td>在Device侧申请的workspace内存地址。</td>
182+ </tr>
183+ <tr>
184+ <td>workspaceSize</td>
185+ <td>输入</td>
186+ <td>在Device侧申请的workspace大小，由第一段接口aclnnCompressorGradGetWorkspaceSize获取。</td>
187+ </tr>
188+ <tr>
189+ <td>executor</td>
190+ <td>输入</td>
191+ <td>op执行器，包含了算子计算流程。</td>
192+ </tr>
193+ <tr>
194+ <td>stream</td>
195+ <td>输入</td>
196+ <td>指定执行任务的Stream。</td>
197+ </tr>
198+ </tbody>
199+ </table>
200+ 
201+- **返回值**
202+ 
203+ aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn_return_code.md)。
204+ 
205+## 约束说明
206+ 
C
Cchenjiao15 天前

缺少确定性说明

likedislike
207+- x参数维度含义：B（Batch Size）表示输入样本批量大小、S（Sequence Length）表示输入样本序列长度、H（Head Size）表示hidden层的大小、D（Head Dim）表示hidden层的最小单元大小、T表示所有Batch输入样本序列长度的累加和。
208+- 输入shape限制：
209+ - wkv支持输入shape[coff* D,H]
210+ - wgate支持输入shape[coff* D,H]
211+ - softmaxScore支持输入shape：BS合轴时为[min(T,T//cmp_ratio+B), coff*cmp_ratio, D]；BS非合轴时为[B,ceil(S/cmp_ratio),coff*cmp_ratio,D]
212+ - kv支持输入shape：同softmaxScore
213+ - dCmpKv支持输入shape：BS合轴时为[min(T,T//cmp_ratio+B),D]；BS非合轴时为[B,ceil(S/cmp_ratio),D]
214+ - dX支持输出shape：与x相同，BS合轴时为[T,H]、BS非合轴时为[B,S,H]
215+ - dWkv、dWgate支持输出shape[coff* D,H]
216+ - dApe支持输出shape[cmp_ratio,coff* D]
217+ - startPos支持输入shape[B,]
218+ - 若x的维度采用BS合轴，即x的输入shape为[T,H]
219+ - cuSeqlens输入shape必须为[B+1,]。该参数中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值，且第一位必须为0。
220+ - seqused，支持输入shape[B,]，要求每个Batch的有效token数要求小于等于对应Sequence Length长度，即seqused[n] <= cu\_seqlens[n+1] - cu\_seqlens[n]，且不小于0。
221+ - 若x的维度不采用BS合轴，即x的输入shape为[B,S,H]
222+ - cuSeqlens，参数必须为空。
223+ - seqused，支持输入shape[B,]，要求每个Batch的有效token数要求小于等于对应Sequence Length长度，即要求seqused[n] <= S，且不小于0。
224+- 输入值域限制：
225+ - 该接口支持B、S泛化，且存在如下场景限制：
226+ - **不支持B、S、T为0的空Tensor**：与正向Compressor不同，CompressorGrad所有输入/输出均不支持空Tensor，shapeSize必须大于0。
227+ - 部分长序列场景下，如果计算量过大可能会导致出现超过NPU内存的报错，注：这里计算量会受x输入shape的影响，值越大计算量越大。
228+- 输入属性限制：
229+ - 支持D为128/512。
230+ - 支持H为1K~10K，512对齐。
231+ - 支持coff为1/2。
232+ - 支持cmp\_ratio为2~128。
233+ 
234+## 调用示例
235+ 
236+ 无
Aattention/compressor_grad/op_host/CMakeLists.txt+52-0
@@ -0,0 +1,52 @@
1+# This program is free software, you can redistribute it and/or modify.
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This file is a part of the CANN Open Software.
4+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
7+# See LICENSE in the root of the software repository for the full text of the License.
8+# ======================================================================================================================
9+require_pypto_pro(compressor_grad)
10+add_op_to_compiled_list()
11+ 
12+if (BUILD_OPEN_PROJECT)
13+ target_sources(op_host_aclnn PRIVATE
14+ compressor_grad_def.cpp
15+ )
16+ enable_pypto_kernel(compressor_grad)
17+endif()
18+ 
19+set(CONDITION_UNIT ${ASCEND_COMPUTE_UNIT})
20+if("${CONDITION_UNIT}" STREQUAL "ascend950")
21+ # PyPTO kernel: runs codegen at configure time and exposes PYPTO_GEN_DIR for the host tiling -I below.
22+ # Must be called BEFORE add_modules_sources(...) (which reads PYPTO_GEN_DIR).
23+ add_ops_compile_options(
24+ OP_NAME CompressorGrad
25+ COMPUTE_UNIT Ascend950PR_9599
26+ OPTIONS -DENABLE_CV_COMM_VIA_SSBUF=true
27+ )
28+endif()
29+ 
30+if (BUILD_OPS_RTY_KERNEL)
31+ add_ops_compile_options(
32+ OP_NAME CompressorGrad
33+ OPTIONS --cce-auto-sync=off
34+ -Wno-deprecated-declarations
35+ -Werror
36+ )
37+else()
38+ add_ops_compile_options(
39+ OP_NAME CompressorGrad
40+ OPTIONS --cce-auto-sync=off
41+ -Wno-deprecated-declarations
42+ )
43+endif()
44+ 
45+if(NOT BUILD_OPS_RTY_KERNEL)
46+ add_modules_sources(OPTYPE compressor_grad ACLNNTYPE aclnn)
47+endif()
48+ 
49+if(NOT BUILD_OPS_RTY_KERNEL)
50+ add_tiling_modules()
51+endif()
52+ 
Aattention/compressor_grad/op_host/compressor_grad_def.cpp+97-0
@@ -0,0 +1,97 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include "register/op_def_registry.h"
12+ 
13+namespace ops {
14+class CompressorGrad : public OpDef {
15+public:
16+ static constexpr uint32_t CMP_RATIO_VALUE = 4;
17+ static constexpr uint32_t COFF_VALUE = 1;
18+ 
19+ explicit CompressorGrad(const char *name) : OpDef(name)
20+ {
21+ this->Input("x")
22+ .ParamType(REQUIRED)
23+ .DataType({ge::DT_BF16, ge::DT_FLOAT16})
24+ .FormatList({ge::FORMAT_ND})
25+ .AutoContiguous();
26+ this->Input("wkv")
27+ .ParamType(REQUIRED)
28+ .DataType({ge::DT_BF16, ge::DT_FLOAT16})
29+ .FormatList({ge::FORMAT_ND})
30+ .AutoContiguous();
31+ this->Input("wgate")
32+ .ParamType(REQUIRED)
33+ .DataType({ge::DT_BF16, ge::DT_FLOAT16})
34+ .FormatList({ge::FORMAT_ND})
35+ .AutoContiguous();
36+ this->Input("d_cmp_kv")
37+ .ParamType(REQUIRED)
38+ .DataType({ge::DT_BF16, ge::DT_FLOAT16})
39+ .FormatList({ge::FORMAT_ND})
40+ .AutoContiguous();
41+ this->Input("softmax_score")
42+ .ParamType(REQUIRED)
43+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
44+ .FormatList({ge::FORMAT_ND})
45+ .AutoContiguous();
46+ this->Input("kv")
47+ .ParamType(REQUIRED)
48+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
49+ .FormatList({ge::FORMAT_ND})
50+ .AutoContiguous();
51+ this->Input("cu_seqlens")
52+ .ParamType(OPTIONAL)
53+ .DataTypeList({ge::DT_INT32})
54+ .FormatList({ge::FORMAT_ND})
55+ .AutoContiguous();
56+ this->Input("seqused")
57+ .ParamType(OPTIONAL)
58+ .DataTypeList({ge::DT_INT32})
59+ .FormatList({ge::FORMAT_ND})
60+ .AutoContiguous();
61+ this->Input("start_pos")
62+ .ParamType(OPTIONAL)
63+ .DataTypeList({ge::DT_INT32})
64+ .FormatList({ge::FORMAT_ND})
65+ .AutoContiguous();
66+ this->Output("d_x")
67+ .ParamType(REQUIRED)
68+ .DataType({ge::DT_BF16, ge::DT_FLOAT16})
69+ .FormatList({ge::FORMAT_ND});
70+ this->Output("d_wkv")
71+ .ParamType(REQUIRED)
72+ .DataType({ge::DT_BF16, ge::DT_FLOAT16})
73+ .FormatList({ge::FORMAT_ND});
74+ this->Output("d_wgate")
75+ .ParamType(REQUIRED)
76+ .DataType({ge::DT_BF16, ge::DT_FLOAT16})
77+ .FormatList({ge::FORMAT_ND});
78+ this->Output("d_ape")
79+ .ParamType(REQUIRED)
80+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
81+ .FormatList({ge::FORMAT_ND});
82+ 
83+ this->Attr("cmp_ratio").AttrType(REQUIRED).Int(CMP_RATIO_VALUE);
84+ this->Attr("coff").AttrType(OPTIONAL).Int(COFF_VALUE);
85+ OpAICoreConfig aicore_config;
86+ aicore_config.DynamicCompileStaticFlag(true)
87+ .DynamicFormatFlag(true)
88+ .DynamicRankSupportFlag(true)
89+ .DynamicShapeSupportFlag(true)
90+ .NeedCheckSupportFlag(false)
91+ .PrecisionReduceFlag(true)
92+ .ExtendCfgInfo("aclnnSupport.value", "support_aclnn"); // set value of aclnn support
93+ this->AICore().AddConfig("ascend950", aicore_config);
94+ }
95+};
96+OP_ADD(CompressorGrad, optiling::CompressorGradCompileInfo);
97+} // namespace ops
Aattention/compressor_grad/op_host/compressor_grad_infershape.cpp+193-0
@@ -0,0 +1,193 @@
1+/* *
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+  */
10+ 
11+/*!
12+ * \file compressor_grad_infershape.cpp
13+ * \brief
14+ */
15+ 
16+#include <graph/utils/type_utils.h>
17+#include <register/op_impl_registry.h>
18+#include "err/ops_err.h"
19+#include "log/log.h"
20+ 
21+using namespace ge;
22+ 
23+namespace ops {
24+ // INPUT
25+ constexpr uint32_t TOKEN_X_INPUT_INDEX = 0;
26+ constexpr uint32_t WEIGHT_KV_INPUT_INDEX = 1;
27+ constexpr uint32_t WEIGHT_WGATE_INPUT_INDEX = 2;
28+ constexpr uint32_t D_CMP_KV_INPUT_INDEX = 3;
29+ constexpr uint32_t SOFTMAX_SCORE_INPUT_INDEX = 4;
30+ constexpr uint32_t KV_INPUT_INDEX = 5;
31+ 
32+ // INPUT(OPTION)
33+ constexpr uint32_t CU_SEQ_LEN_INPUT_INDEX = 6;
34+ constexpr uint32_t SEQ_USED_INPUT_INDEX = 7;
35+ constexpr uint32_t START_POS_INPUT_INDEX = 8;
36+ 
37+ // ATTR
38+ constexpr uint32_t CMP_RATIO_ATTR_INDEX = 0;
39+ constexpr uint32_t COFF_ATTR_INDEX = 1;
40+ 
41+ // OUTPUT
42+ constexpr uint32_t D_X_OUTPUT_INDEX = 0;
43+ constexpr uint32_t D_WKV_OUTPUT_INDEX = 1;
44+ constexpr uint32_t D_WGATE_OUTPUT_INDEX = 2;
45+ constexpr uint32_t D_APE_OUTPUT_INDEX = 3;
46+ 
47+ // ATTR DEFAULT VALUE
48+ constexpr uint32_t CMP_RATIO_VALUE = 4;
49+ constexpr uint32_t COFF_VALUE = 1;
50+ 
51+struct CompressorGradProtoShapeParam {
52+ bool isBsMerge { false };
53+ int64_t B { 0 };
54+ int64_t T { 0 };
55+ int64_t S { 0 };
56+ int64_t Sr { 0 };
57+ int64_t H { 0 };
58+ int64_t D { 0 };
59+};
60+ 
61+// tmp
62+constexpr uint32_t DIM_NUM_1 = 1;
63+constexpr uint32_t DIM_NUM_2 = 2;
64+constexpr uint32_t DIM_NUM_3 = 3;
65+constexpr uint32_t DIM_NUM_4 = 4;
66+constexpr uint32_t DIM_INDEX_0 = 0;
67+constexpr uint32_t DIM_INDEX_1 = 1;
68+constexpr uint32_t DIM_INDEX_2 = 2;
69+constexpr uint32_t DIM_INDEX_3 = 3;
70+ 
71+ge::graphStatus GetCompressorGradShapeDim(const gert::InferShapeContext* context,
72+ CompressorGradProtoShapeParam &shapeParam)
73+{
74+ auto xShape = context->GetRequiredInputShape(TOKEN_X_INPUT_INDEX); // (B, S, H) | (T, H)
75+ OP_CHECK_NULL_WITH_CONTEXT(context, xShape);
76+ auto wkvShape = context->GetRequiredInputShape(WEIGHT_KV_INPUT_INDEX); // (coff * D, H)
77+ OP_CHECK_NULL_WITH_CONTEXT(context, wkvShape);
78+ auto wgateShape = context->GetRequiredInputShape(WEIGHT_WGATE_INPUT_INDEX); // (coff * D, H)
79+ OP_CHECK_NULL_WITH_CONTEXT(context, wgateShape);
80+ 
81+ auto attr = context->GetAttrs();
82+ OP_CHECK_NULL_WITH_CONTEXT(context, attr);
83+ const int64_t *cmpRatioPtr = attr->GetAttrPointer<int64_t>(CMP_RATIO_ATTR_INDEX);
84+ int64_t cmpRatio = (cmpRatioPtr != nullptr) ? *cmpRatioPtr : CMP_RATIO_VALUE;
85+ const int64_t *coffPtr = attr->GetAttrPointer<int64_t>(COFF_ATTR_INDEX);
86+ int64_t coff = (coffPtr != nullptr) ? *coffPtr : COFF_VALUE;
87+ 
88+ if (xShape->GetDimNum() == DIM_NUM_3) { // BS
89+ shapeParam.isBsMerge = false;
90+ shapeParam.B = xShape->GetDim(DIM_INDEX_0);
91+ shapeParam.S = xShape->GetDim(DIM_INDEX_1);
92+ shapeParam.H = xShape->GetDim(DIM_INDEX_2);
93+ shapeParam.T = shapeParam.B * shapeParam.S;
94+ } else { // T
95+ shapeParam.isBsMerge = true;
96+ auto cuSeqlensShape = context->GetOptionalInputShape(CU_SEQ_LEN_INPUT_INDEX);
97+ shapeParam.T = xShape->GetDim(DIM_INDEX_0);
98+ shapeParam.H = xShape->GetDim(DIM_INDEX_1);
99+ }
100+ 
101+ shapeParam.D = wkvShape->GetDim(DIM_INDEX_0) / coff;
102+ 
103+ return GRAPH_SUCCESS;
104+}
105+ 
106+ge::graphStatus SetCompressorGradShapeDim(const CompressorGradProtoShapeParam &shapeParam,
107+ gert::InferShapeContext* context)
108+{
109+ // d_x: 与 x 同形状 (B, S, H) | (T, H)
110+ auto dxShape = context->GetOutputShape(D_X_OUTPUT_INDEX);
111+ OP_CHECK_NULL_WITH_CONTEXT(context, dxShape);
112+ 
113+ // d_wkv: (coff * D, H)
114+ auto dWkvShape = context->GetOutputShape(D_WKV_OUTPUT_INDEX);
115+ OP_CHECK_NULL_WITH_CONTEXT(context, dWkvShape);
116+ 
117+ auto dWgateShape = context->GetOutputShape(D_WGATE_OUTPUT_INDEX);
118+ OP_CHECK_NULL_WITH_CONTEXT(context, dWgateShape);
119+ 
120+ auto dApeShape = context->GetOutputShape(D_APE_OUTPUT_INDEX);
121+ OP_CHECK_NULL_WITH_CONTEXT(context, dApeShape);
122+ 
123+ auto attr = context->GetAttrs();
124+ OP_CHECK_NULL_WITH_CONTEXT(context, attr);
125+ const int64_t *cmpRatioPtr = attr->GetAttrPointer<int64_t>(CMP_RATIO_ATTR_INDEX);
126+ int64_t cmpRatio = (cmpRatioPtr != nullptr) ? *cmpRatioPtr : CMP_RATIO_VALUE;
127+ const int64_t *coffPtr = attr->GetAttrPointer<int64_t>(COFF_ATTR_INDEX);
128+ int64_t coff = (coffPtr != nullptr) ? *coffPtr : COFF_VALUE;
129+ // Set output shape
130+ if (!shapeParam.isBsMerge) {
131+ dxShape->SetDimNum(DIM_NUM_3); // (B, S, H)
132+ dxShape->SetDim(DIM_INDEX_0, shapeParam.B);
133+ dxShape->SetDim(DIM_INDEX_1, shapeParam.S);
134+ dxShape->SetDim(DIM_INDEX_2, shapeParam.H);
135+ } else {
136+ dxShape->SetDimNum(DIM_NUM_2); // (T, H)
137+ dxShape->SetDim(DIM_INDEX_0, shapeParam.T);
138+ dxShape->SetDim(DIM_INDEX_1, shapeParam.H);
139+ }
140+ 
141+ dWkvShape->SetDimNum(DIM_NUM_2); // (coff*D, H)
142+ dWkvShape->SetDim(DIM_INDEX_0, shapeParam.D * coff); // (coff*D, H)
143+ dWkvShape->SetDim(DIM_INDEX_1, shapeParam.H); // (coff*D, H)
144+ 
145+ dWgateShape->SetDimNum(DIM_NUM_2); // (coff*D, H)
146+ dWgateShape->SetDim(DIM_INDEX_0, shapeParam.D * coff); // (coff*D, H)
147+ dWgateShape->SetDim(DIM_INDEX_1, shapeParam.H); // (coff*D, H)
148+ 
149+ dApeShape->SetDimNum(DIM_NUM_2); // (cmpRatio, coff*D)
150+ dApeShape->SetDim(DIM_INDEX_0, cmpRatio); // (cmpRatio, coff*D)
151+ dApeShape->SetDim(DIM_INDEX_1, shapeParam.D * coff); // (cmpRatio, coff*D)
152+ 
153+ return GRAPH_SUCCESS;
154+}
155+ 
156+ge::graphStatus InferDataTypeCompressorGrad(gert::InferDataTypeContext* context)
157+{
158+ OP_CHECK_IF(context == nullptr,
159+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON("CompressorGrad", "context", "is nullptr"),
160+ return ge::GRAPH_FAILED);
161+ OP_LOGI(context->GetNodeName(), "Enter CompressorGrad inferDataType impl.");
162+ 
163+ context->SetOutputDataType(D_X_OUTPUT_INDEX, context->GetRequiredInputDataType(TOKEN_X_INPUT_INDEX));
164+ context->SetOutputDataType(D_WKV_OUTPUT_INDEX, context->GetRequiredInputDataType(TOKEN_X_INPUT_INDEX));
165+ context->SetOutputDataType(D_WGATE_OUTPUT_INDEX, context->GetRequiredInputDataType(TOKEN_X_INPUT_INDEX));
166+ context->SetOutputDataType(D_APE_OUTPUT_INDEX, context->GetRequiredInputDataType(SOFTMAX_SCORE_INPUT_INDEX));
167+ 
168+ return GRAPH_SUCCESS;
169+}
170+ 
171+ge::graphStatus InferShapeCompressorGrad(gert::InferShapeContext* context)
172+{
173+ OP_CHECK_IF(context == nullptr,
174+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON("CompressorGrad", "context", "is nullptr"),
175+ return ge::GRAPH_FAILED);
176+ OP_LOGI(context->GetNodeName(), "Enter CompressorGrad infershape impl.");
177+ 
178+ CompressorGradProtoShapeParam shapeParam {};
179+ auto apiRet = GetCompressorGradShapeDim(context, shapeParam);
180+ OP_CHECK_IF((apiRet != GRAPH_SUCCESS),
181+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context->GetNodeName(), "context", "get input shape failed"),
182+ return ge::GRAPH_FAILED);
183+ 
184+ apiRet = SetCompressorGradShapeDim(shapeParam, context);
185+ OP_CHECK_IF((apiRet != GRAPH_SUCCESS),
186+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context->GetNodeName(), "context", "set output shape failed"),
187+ return ge::GRAPH_FAILED);
188+ 
189+ return GRAPH_SUCCESS;
190+}
191+ 
192+IMPL_OP_INFERSHAPE(CompressorGrad).InferShape(InferShapeCompressorGrad).InferDataType(InferDataTypeCompressorGrad);
193+} // namespace ops
Aattention/compressor_grad/op_host/compressor_grad_tiling.cpp+1097-0
@@ -0,0 +1,1097 @@
1+/* *
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+  */
10+ 
11+/*!
12+ * \file compressor_grad_tiling.cpp
13+ * \brief
14+ */
15+ 
16+#include <functional>
17+#include <algorithm>
18+#include <unordered_map>
19+#include <graph/utils/type_utils.h>
20+#include "log/log.h"
21+#include "register/op_def_registry.h"
22+#include "compressor_grad_tiling.h"
23+ 
24+using namespace ge;
25+using namespace AscendC;
26+namespace optiling {
27+ 
28+void ConvertRequiredParams(gert::TilingContext &context, CompressorGradContext &compressorGradContext)
29+{
30+ compressorGradContext.x.desc = context.GetRequiredInputDesc(TOKEN_X_INPUT_INDEX);
31+ compressorGradContext.x.shape = context.GetRequiredInputShape(TOKEN_X_INPUT_INDEX);
32+ compressorGradContext.wkv.desc = context.GetRequiredInputDesc(WEIGHT_KV_INPUT_INDEX);
33+ compressorGradContext.wkv.shape = context.GetRequiredInputShape(WEIGHT_KV_INPUT_INDEX);
34+ compressorGradContext.wgate.desc = context.GetRequiredInputDesc(WEIGHT_WGATE_INPUT_INDEX);
35+ compressorGradContext.wgate.shape = context.GetRequiredInputShape(WEIGHT_WGATE_INPUT_INDEX);
36+ compressorGradContext.dCmpKv.desc = context.GetRequiredInputDesc(D_CMP_KV_INPUT_INDEX);
37+ compressorGradContext.dCmpKv.shape = context.GetRequiredInputShape(D_CMP_KV_INPUT_INDEX);
38+ compressorGradContext.softmaxScore.desc = context.GetRequiredInputDesc(SOFTMAX_SCORE_INPUT_INDEX);
39+ compressorGradContext.softmaxScore.shape = context.GetRequiredInputShape(SOFTMAX_SCORE_INPUT_INDEX);
40+ compressorGradContext.kv.desc = context.GetRequiredInputDesc(KV_INPUT_INDEX);
41+ compressorGradContext.kv.shape = context.GetRequiredInputShape(KV_INPUT_INDEX);
42+ 
43+ compressorGradContext.dX.desc = context.GetOutputDesc(D_X_OUTPUT_INDEX);
44+ compressorGradContext.dX.shape = context.GetOutputShape(D_X_OUTPUT_INDEX);
45+ compressorGradContext.dWkv.desc = context.GetOutputDesc(D_WKV_OUTPUT_INDEX);
46+ compressorGradContext.dWkv.shape = context.GetOutputShape(D_WKV_OUTPUT_INDEX);
47+ compressorGradContext.dWgate.desc = context.GetOutputDesc(D_WGATE_OUTPUT_INDEX);
48+ compressorGradContext.dWgate.shape = context.GetOutputShape(D_WGATE_OUTPUT_INDEX);
49+ compressorGradContext.dApe.desc = context.GetOutputDesc(D_APE_OUTPUT_INDEX);
50+ compressorGradContext.dApe.shape = context.GetOutputShape(D_APE_OUTPUT_INDEX);
51+ 
52+ compressorGradContext.dtype = compressorGradContext.x.desc->GetDataType();
53+ auto xDimNum = compressorGradContext.x.shape->GetStorageShape().GetDimNum();
54+ if (xDimNum == COMPRESSOR_GRAD_DIM_NUM_3) {
55+ compressorGradContext.layout = LayoutType::LAYOUT_BSH;
56+ } else if (xDimNum == COMPRESSOR_GRAD_DIM_NUM_2) {
57+ compressorGradContext.layout = LayoutType::LAYOUT_TH;
58+ }
59+}
60+ 
61+void ConvertOptionalParams(gert::TilingContext &context, CompressorGradContext &compressorGradContext)
62+{
63+ compressorGradContext.cuSeqlens.desc = context.GetOptionalInputDesc(CU_SEQ_LEN_INPUT_INDEX);
64+ compressorGradContext.cuSeqlens.shape = context.GetOptionalInputShape(CU_SEQ_LEN_INPUT_INDEX);
65+ compressorGradContext.seqUsed.desc = context.GetOptionalInputDesc(SEQ_USED_INPUT_INDEX);
66+ compressorGradContext.seqUsed.shape = context.GetOptionalInputShape(SEQ_USED_INPUT_INDEX);
67+ compressorGradContext.startPos.desc = context.GetOptionalInputDesc(START_POS_INPUT_INDEX);
68+ compressorGradContext.startPos.shape = context.GetOptionalInputShape(START_POS_INPUT_INDEX);
69+}
70+ 
71+ge::graphStatus ConvertContext(gert::TilingContext &context, CompressorGradContext &compressorGradContext)
72+{
73+ if (context.GetNodeName() == nullptr) {
74+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON("CompressorGrad", "opName", "got from TilingContext is nullptr");
75+ return ge::GRAPH_FAILED;
76+ }
77+ 
78+ compressorGradContext.opName = context.GetNodeName();
79+ compressorGradContext.opType = context.GetNodeType();
80+ ConvertRequiredParams(context, compressorGradContext);
81+ ConvertOptionalParams(context, compressorGradContext);
82+ 
83+ auto attrs = context.GetAttrs();
84+ OP_CHECK_IF(attrs == nullptr,
85+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.GetNodeName(), "attrs", "got from ge is nullptr"),
86+ return ge::GRAPH_FAILED);
87+ compressorGradContext.coff = attrs->GetAttrPointer<uint32_t>(COFF_ATTR_INDEX);
88+ compressorGradContext.cmpRatio = attrs->GetAttrPointer<uint32_t>(CMP_RATIO_ATTR_INDEX);
89+ 
90+ OP_CHECK_IF(context.GetWorkspaceSizes(1) == nullptr,
91+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.GetNodeName(), "workSpaceSize",
92+ "got from ge is nullptr"),
93+ return ge::GRAPH_FAILED);
94+ compressorGradContext.workSpaces = context.GetWorkspaceSizes(1);
95+ return ge::GRAPH_SUCCESS;
96+}
97+ 
98+ge::graphStatus SetBaseInfo(CompressorGradContext &compressorGradContext,
99+ CompressorGradBaseParams &compressorGradBaseParams, uint32_t aicNum)
100+{
101+ if (compressorGradContext.x.shape->GetStorageShape().GetDimNum() == COMPRESSOR_GRAD_DIM_NUM_3) {
102+ compressorGradBaseParams.batchSize =
103+ compressorGradContext.x.shape->GetStorageShape().GetDim(COMPRESSOR_GRAD_DIM_INDEX_0);
104+ compressorGradBaseParams.seqSize =
105+ compressorGradContext.x.shape->GetStorageShape().GetDim(COMPRESSOR_GRAD_DIM_INDEX_1);
106+ compressorGradBaseParams.hiddenSize =
107+ compressorGradContext.x.shape->GetStorageShape().GetDim(COMPRESSOR_GRAD_DIM_INDEX_2);
108+ compressorGradBaseParams.tokenSize =
109+ compressorGradBaseParams.batchSize * compressorGradBaseParams.seqSize;
110+ } else {
111+ compressorGradBaseParams.batchSize =
112+ compressorGradContext.cuSeqlens.shape->GetStorageShape().GetDim(COMPRESSOR_GRAD_DIM_INDEX_0) - 1;
113+ compressorGradBaseParams.tokenSize =
114+ compressorGradContext.x.shape->GetStorageShape().GetDim(COMPRESSOR_GRAD_DIM_INDEX_0);
115+ compressorGradBaseParams.hiddenSize =
116+ compressorGradContext.x.shape->GetStorageShape().GetDim(COMPRESSOR_GRAD_DIM_INDEX_1);
117+ }
118+ 
119+ uint8_t coff = compressorGradContext.coff == nullptr ? COFF_VALUE :
120+ static_cast<uint8_t>(*compressorGradContext.coff);
121+ compressorGradBaseParams.headDim =
122+ compressorGradContext.wkv.shape->GetStorageShape().GetDim(COMPRESSOR_GRAD_DIM_INDEX_0) / coff;
123+ compressorGradBaseParams.featureDim =
124+ compressorGradContext.wkv.shape->GetStorageShape().GetDim(COMPRESSOR_GRAD_DIM_INDEX_0);
125+ compressorGradBaseParams.cmpRatio = static_cast<uint32_t>(*compressorGradContext.cmpRatio);
126+ compressorGradBaseParams.nSize = 2; // 预留（当前未参与 tiling 决策）
127+ compressorGradBaseParams.usedCoreNum = aicNum;
128+ OP_LOGI(compressorGradContext.opName, "[TILING] bSize:%u tSize:%u cmpRatio:%u coff:%u",
129+ compressorGradBaseParams.batchSize, compressorGradBaseParams.tokenSize,
130+ compressorGradBaseParams.cmpRatio, coff);
131+ 
132+ return ge::GRAPH_SUCCESS;
133+}
134+ 
135+ge::graphStatus CalcWorkSpace(CompressorGradContext &compressorGradContext,
136+ CompressorGradBaseParams &compressorGradBaseParams,
137+ CompressorGradWorkspaceParams &compressorGradWorkspaceParams,
138+ size_t &workspaceSize, size_t libapiSize, uint32_t aicNum)
139+{
140+ constexpr uint32_t MM1_RES_ELEM_SIZE = 4; // 4: fp32
141+ constexpr uint32_t M_BASE_SIZE = 128;
142+
143+ uint8_t coff = compressorGradContext.coff == nullptr ? COFF_VALUE :
144+ static_cast<uint8_t>(*compressorGradContext.coff);
145+ uint32_t cmpRatio = compressorGradBaseParams.cmpRatio;
146+ uint32_t headDim = compressorGradBaseParams.headDim;
147+ uint32_t hiddenSize = compressorGradBaseParams.hiddenSize;
148+ uint32_t cubeCoreNum = aicNum;
149+ uint32_t groupSize = headDim / 128; // 与 kernel groupSize = headDim // D_BASE_SIZE 一致
150+ uint32_t groupNum = cubeCoreNum / groupSize;
151+ uint32_t cmpSize = coff * cmpRatio * headDim;
152+ uint32_t totalHeadDim = coff * headDim;
153+ uint32_t coffCoef = 2 / coff;
154+ uint32_t dealScNum = 128 / cmpRatio;
155+ uint32_t groupDealScNum = dealScNum * coffCoef;
156+ uint32_t groupRowStride = groupDealScNum * cmpRatio + (coff - 1) * cmpRatio; // 与 kernel 一致
157+ uint32_t dbRatio = 2;
158+ 
159+ // 与 kernel 的 workspace 指针链逐分区对齐（元素数，FP32）:
160+ // ape / dWkv / dWgate 单缓冲；dX / x / dXCache 按 dbRatio=2 双缓冲
161+ uint64_t apeWorkSpaceSize = static_cast<uint64_t>(groupNum) * cmpSize * coffCoef;
162+ uint64_t dXWorkSpaceSize = static_cast<uint64_t>(dbRatio) * cubeCoreNum * (M_BASE_SIZE * 2) * hiddenSize;
163+ uint64_t dWeightWorkSpaceSize = static_cast<uint64_t>(groupNum) * totalHeadDim * hiddenSize;
164+ // dWkv / dWGate 各占一份 dWeightWorkSpaceSize
165+ uint64_t xWorkSpaceSize = static_cast<uint64_t>(dbRatio) * groupNum * groupRowStride * hiddenSize;
166+ uint64_t dXCacheWorkSpaceSize = static_cast<uint64_t>(dbRatio) * cmpRatio * hiddenSize;
167+ 
168+ workspaceSize = libapiSize;
169+ workspaceSize += (apeWorkSpaceSize + dXWorkSpaceSize + dWeightWorkSpaceSize * 2 +
170+ xWorkSpaceSize + dXCacheWorkSpaceSize) * MM1_RES_ELEM_SIZE;
171+ 
172+ if (compressorGradContext.workSpaces) {
173+ compressorGradContext.workSpaces[0] = workspaceSize;
174+ }
175+ 
176+ OP_LOGI(compressorGradContext.opName,
177+ "Tiling info: workspaceSize = %zu (ape=%llu dx=%llu dw=%llu x=%llu dxcache=%llu)",
178+ workspaceSize, apeWorkSpaceSize, dXWorkSpaceSize,
179+ dWeightWorkSpaceSize * 2, xWorkSpaceSize, dXCacheWorkSpaceSize);
180+ return ge::GRAPH_SUCCESS;
181+}
182+ 
183+template <typename T>
184+std::string to_string(const T &value)
185+{
186+ if (std::is_same_v<T, bool>) {
187+ return value ? "true" : "false";
188+ } else {
189+ return std::to_string(value);
190+ }
191+}
192+ 
193+ge::graphStatus GenTilingKey(CompressorGradContext &compressorGradContext)
194+{
195+ // 0:BF16, 1:FP16
196+ uint8_t dtype = 0;
197+ // 0: BSH 1:TH
198+ uint8_t layout = 0;
199+ 
200+ auto xDtype = compressorGradContext.x.desc->GetDataType();
201+ if (xDtype == ge::DT_BF16) {
202+ dtype = 0;
203+ } else if (xDtype == ge::DT_FLOAT16) {
204+ dtype = 1;
205+ }
206+ auto xDimNum = compressorGradContext.x.shape->GetStorageShape().GetDimNum();
207+ if (xDimNum == COMPRESSOR_GRAD_DIM_NUM_3) {
208+ layout = 0;
209+ } else {
210+ layout = 1;
211+ }
212+ 
213+ uint8_t coff = compressorGradContext.coff == nullptr ? COFF_VALUE :
214+ static_cast<uint8_t>(*compressorGradContext.coff);
215+ // 通过 ASCENDC 宏编码 tilingKey（force-include 的 codegen 生成头
216+ // CompressorGradTilingKey_tilingkey.h 声明位布局，与 PyPTO 一致：
217+ // Coff(2bit) | Layout(1bit) | DataType(2bit)，UINT 值自动映射为索引）
218+ compressorGradContext.tilingKey = GET_TPL_TILING_KEY(coff, layout, dtype);
219+
220+ OP_LOGI(compressorGradContext.opName,
221+ "CompressorGrad dtype:%hhu layout:%hhu coff:%hhu", dtype, layout, coff);
222+ OP_LOGI(compressorGradContext.opName, "CompressorGrad tilingKey:%lu",
223+ compressorGradContext.tilingKey);
224+ return ge::GRAPH_SUCCESS;
225+}
226+ 
227+template <typename T>
228+void LogErrorNumberSupport(const std::vector<T> &expectNumberList, const T &actualValue,
229+ const std::string &name, const std::string subName, const char *opName)
230+{
231+ std::ostringstream oss;
232+ for (size_t i = 0; i < expectNumberList.size(); ++i) {
233+ oss << to_string(expectNumberList[i]);
234+ if (i < expectNumberList.size() - 1) {
235+ oss << ", ";
236+ }
237+ }
238+ 
239+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName, name, to_string(actualValue),
240+ subName + " only supports " + oss.str());
241+}
242+ 
243+template <typename T>
244+ge::graphStatus CheckFeatureValueSupport(const T *featureValue,
245+ const std::vector<T> &expectFeatureValList,
246+ const std::string &name, const char* opName)
247+{
248+ if (std::find(expectFeatureValList.begin(), expectFeatureValList.end(), *featureValue) ==
249+ expectFeatureValList.end()) {
250+ LogErrorNumberSupport(expectFeatureValList, *featureValue, name, "feature value", opName);
251+ return ge::GRAPH_FAILED;
252+ }
253+ return ge::GRAPH_SUCCESS;
254+}
255+ 
256+template <typename T>
257+ge::graphStatus CheckAttrValueSupportInterval(const T *attrValue, const uint32_t minVal, const uint32_t maxVal,
258+ const std::string &name, const char *opName)
259+{
260+ if (attrValue == nullptr) {
261+ return ge::GRAPH_SUCCESS;
262+ }
263+
264+ std::string attr_value = "attr value";
265+ if (*attrValue < minVal || *attrValue > maxVal) {
266+ std::ostringstream oss;
267+ oss << "[" << minVal << ", " << maxVal << "]";
268+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName, name, std::to_string(*attrValue),
269+ "attr value only supports value in range " + oss.str());
270+ return ge::GRAPH_FAILED;
271+ }
272+ 
273+ return ge::GRAPH_SUCCESS;
274+}
275+ 
276+template <typename T>
277+ge::graphStatus CheckAttrValueSupportList(const T *attrValue, const std::vector<T> &expectAttrValList,
278+ const std::string &name, const char *opName)
279+{
280+ if (attrValue == nullptr) {
281+ return ge::GRAPH_SUCCESS;
282+ }
283+ 
284+ if (std::find(expectAttrValList.begin(), expectAttrValList.end(), *attrValue) == expectAttrValList.end()) {
285+ LogErrorNumberSupport(expectAttrValList, *attrValue, name, "attr value", opName);
286+ return ge::GRAPH_FAILED;
287+ }
288+ 
289+ return ge::GRAPH_SUCCESS;
290+}
291+ 
292+namespace compressor_grad_tiling {
293+std::string LayoutTypeToStr(LayoutType layout)
294+{
295+ switch (layout) {
296+ case LayoutType::LAYOUT_BSH:
297+ return "BSH";
298+ case LayoutType::LAYOUT_TH:
299+ return "TH";
300+ default:
301+ return "UNKNOWN_LAYOUT";
302+ }
303+}
304+} // namespace compressor_grad_tiling
305+ 
306+ge::graphStatus CheckDimNumInLayoutSupport(CompressorGradContext &compressorGradContext, const std::string &layout,
307+ const gert::StorageShape *shape, const std::string &name)
308+{
309+ const auto &dimIt = LAYOUT_DIM_MAP.find(layout);
310+ OP_CHECK_IF(shape->GetStorageShape().GetDimNum() != dimIt->second,
311+ OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(compressorGradContext.opName, name,
312+ std::to_string(shape->GetStorageShape().GetDimNum()),
313+ "when layout is " + layout + ", dimension should be " + std::to_string(dimIt->second)),
314+ return ge::GRAPH_FAILED);
315+ return ge::GRAPH_SUCCESS;
316+}
317+ 
318+void LogErrorDtypeSupport(const std::vector<ge::DataType> &expectDtypeList,
319+ const ge::DataType &actualDtype, const std::string &name, const char *opName)
320+{
321+ std::ostringstream oss;
322+ for (size_t i = 0; i < expectDtypeList.size(); ++i) {
323+ oss << DataTypeToSerialString(expectDtypeList[i]);
324+ if (i < expectDtypeList.size() - 1) {
325+ oss << ", ";
326+ }
327+ }
328+ OP_LOGE_FOR_INVALID_DTYPE(opName, name, DataTypeToSerialString(actualDtype), oss.str());
329+}
330+ 
331+ge::graphStatus CheckDtypeSupport(CompressorGradContext &compressorGradContext, const gert::CompileTimeTensorDesc *desc,
332+ const std::string &name)
333+{
334+ if (desc != nullptr) {
335+ const auto &it = DTYPE_SUPPORT_MAP.find(name);
336+ OP_CHECK_IF(
337+ it == DTYPE_SUPPORT_MAP.end(),
338+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(compressorGradContext.opName, name,
339+ "datatype support list should be specify in DTYPE_SUPPORT_MAP"),
340+ return ge::GRAPH_FAILED);
341+ auto &expectDtypeList = it->second;
342+ OP_CHECK_IF(std::find(expectDtypeList.begin(), expectDtypeList.end(), desc->GetDataType()) ==
343+ expectDtypeList.end(),
344+ LogErrorDtypeSupport(expectDtypeList, desc->GetDataType(),
345+ name, compressorGradContext.opName), return ge::GRAPH_FAILED);
346+ }
347+ return ge::GRAPH_SUCCESS;
348+}
349+ 
350+static std::string DataTypeToSerialString(ge::DataType type)
351+{
352+ const auto it = DATATYPE_TO_STRING_MAP.find(type);
353+ if (it != DATATYPE_TO_STRING_MAP.end()) {
354+ return it->second;
355+ } else {
356+ OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON("CompressorGrad", "datatype", std::to_string(static_cast<int32_t>(type)),
357+ "not support");
358+ return "UNDEFINED";
359+ }
360+}
361+ 
362+ge::graphStatus CheckDimNumSupport(const gert::StorageShape *shape, const std::string &name, const char *opName)
363+{
364+ if (shape == nullptr) {
365+ return ge::GRAPH_SUCCESS;
366+ }
367+ const auto &it = DIM_NUM_MAP.find(name);
368+ OP_CHECK_IF(it == DIM_NUM_MAP.end(),
369+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(opName, name,
370+ "dim number support list should be specify in DIM_NUM_MAP"),
371+ return ge::GRAPH_FAILED);
372+ auto &expectDimNumList = it->second;
373+ OP_CHECK_IF(std::find(expectDimNumList.begin(), expectDimNumList.end(), shape->GetStorageShape().GetDimNum()) ==
374+ expectDimNumList.end(),
375+ LogErrorNumberSupport(expectDimNumList, static_cast<uint32_t>(shape->GetStorageShape().GetDimNum()),
376+ name, "dimension", opName),
377+ return ge::GRAPH_FAILED);
378+ return ge::GRAPH_SUCCESS;
379+}
380+ 
381+ge::graphStatus CheckSingleParaX(CompressorGradContext &compressorGradContext)
382+{
383+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(compressorGradContext, compressorGradContext.x.desc, X_NAME) ||
384+ ge::GRAPH_SUCCESS != CheckDimNumSupport(compressorGradContext.x.shape, X_NAME, compressorGradContext.opName) ||
385+ ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(compressorGradContext,
386+ compressor_grad_tiling::LayoutTypeToStr(compressorGradContext.layout),
387+ compressorGradContext.x.shape, X_NAME)) {
388+ return ge::GRAPH_FAILED;
389+ }
390+ return ge::GRAPH_SUCCESS;
391+}
392+ 
393+ge::graphStatus CheckSingleParaWkv(CompressorGradContext &compressorGradContext)
394+{
395+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(compressorGradContext, compressorGradContext.wkv.desc, WKV_NAME) ||
396+ ge::GRAPH_SUCCESS != CheckDimNumSupport(compressorGradContext.wkv.shape,
397+ WKV_NAME, compressorGradContext.opName)) {
398+ return ge::GRAPH_FAILED;
399+ }
400+ return ge::GRAPH_SUCCESS;
401+}
402+ 
403+ge::graphStatus CheckSingleParaWgate(CompressorGradContext &compressorGradContext)
404+{
405+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(compressorGradContext, compressorGradContext.wgate.desc, WGATE_NAME) ||
406+ ge::GRAPH_SUCCESS != CheckDimNumSupport(compressorGradContext.wgate.shape,
407+ WGATE_NAME, compressorGradContext.opName)) {
408+ return ge::GRAPH_FAILED;
409+ }
410+ return ge::GRAPH_SUCCESS;
411+}
412+ 
413+ge::graphStatus CheckSingleParaDCmpKv(CompressorGradContext &compressorGradContext)
414+{
415+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(compressorGradContext,
416+ compressorGradContext.dCmpKv.desc, D_CMP_KV_NAME) ||
417+ ge::GRAPH_SUCCESS != CheckDimNumSupport(compressorGradContext.dCmpKv.shape,
418+ D_CMP_KV_NAME, compressorGradContext.opName)) {
419+ return ge::GRAPH_FAILED;
420+ }
421+ return ge::GRAPH_SUCCESS;
422+}
423+ 
424+ge::graphStatus CheckSingleParaSoftmaxScore(CompressorGradContext &compressorGradContext)
425+{
426+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(compressorGradContext,
427+ compressorGradContext.softmaxScore.desc,
428+ SOFTMAX_SCORE_NAME) ||
429+ ge::GRAPH_SUCCESS != CheckDimNumSupport(compressorGradContext.softmaxScore.shape,
430+ SOFTMAX_SCORE_NAME,
431+ compressorGradContext.opName)) {
432+ return ge::GRAPH_FAILED;
433+ }
434+ return ge::GRAPH_SUCCESS;
435+}
436+ 
437+ 
438+ge::graphStatus CheckSingleParaKV(CompressorGradContext &compressorGradContext)
439+{
440+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(compressorGradContext, compressorGradContext.kv.desc, KV_NAME) ||
441+ ge::GRAPH_SUCCESS != CheckDimNumSupport(compressorGradContext.kv.shape,
442+ KV_NAME, compressorGradContext.opName)) {
443+ return ge::GRAPH_FAILED;
444+ }
445+ return ge::GRAPH_SUCCESS;
446+}
447+ 
448+ge::graphStatus CheckSingleParaCuSeqlens(CompressorGradContext &compressorGradContext)
449+{
450+ if (compressorGradContext.cuSeqlens.desc == nullptr) {
451+ return ge::GRAPH_SUCCESS;
452+ }
453+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(compressorGradContext,
454+ compressorGradContext.cuSeqlens.desc,
455+ CU_SEQLENS_NAME) ||
456+ ge::GRAPH_SUCCESS != CheckDimNumSupport(compressorGradContext.cuSeqlens.shape,
457+ CU_SEQLENS_NAME,
458+ compressorGradContext.opName)) {
459+ return ge::GRAPH_FAILED;
460+ }
461+ return ge::GRAPH_SUCCESS;
462+}
463+ 
464+ge::graphStatus CheckSingleParaSeqused(CompressorGradContext &compressorGradContext)
465+{
466+ if (compressorGradContext.seqUsed.desc == nullptr) {
467+ return ge::GRAPH_SUCCESS;
468+ }
469+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(compressorGradContext,
470+ compressorGradContext.seqUsed.desc,
471+ SEQUSED_NAME) ||
472+ ge::GRAPH_SUCCESS != CheckDimNumSupport(compressorGradContext.seqUsed.shape,
473+ SEQUSED_NAME,
474+ compressorGradContext.opName)) {
475+ return ge::GRAPH_FAILED;
476+ }
477+ return ge::GRAPH_SUCCESS;
478+}
479+ 
480+ge::graphStatus CheckSingleParaStartPos(CompressorGradContext &compressorGradContext)
481+{
482+ if (compressorGradContext.startPos.desc == nullptr) {
483+ return ge::GRAPH_SUCCESS;
484+ }
485+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(compressorGradContext,
486+ compressorGradContext.startPos.desc,
487+ START_POS_NAME) ||
488+ ge::GRAPH_SUCCESS != CheckDimNumSupport(compressorGradContext.startPos.shape,
489+ START_POS_NAME,
490+ compressorGradContext.opName)) {
491+ return ge::GRAPH_FAILED;
492+ }
493+ return ge::GRAPH_SUCCESS;
494+}
495+ 
496+ge::graphStatus CheckSingleParaDX(CompressorGradContext &compressorGradContext)
497+{
498+ if (compressorGradContext.dX.desc == nullptr) {
499+ return ge::GRAPH_SUCCESS;
500+ }
501+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(compressorGradContext,
502+ compressorGradContext.dX.desc,
503+ D_X_NAME) ||
504+ ge::GRAPH_SUCCESS != CheckDimNumSupport(compressorGradContext.dX.shape,
505+ D_X_NAME,
506+ compressorGradContext.opName)) {
507+ return ge::GRAPH_FAILED;
508+ }
509+ return ge::GRAPH_SUCCESS;
510+}
511+ 
512+ge::graphStatus CheckSingleParaDWkv(CompressorGradContext &compressorGradContext)
513+{
514+ if (compressorGradContext.dWkv.desc == nullptr) {
515+ return ge::GRAPH_SUCCESS;
516+ }
517+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(compressorGradContext,
518+ compressorGradContext.dWkv.desc,
519+ D_WKV_NAME) ||
520+ ge::GRAPH_SUCCESS != CheckDimNumSupport(compressorGradContext.dWkv.shape,
521+ D_WKV_NAME,
522+ compressorGradContext.opName)) {
523+ return ge::GRAPH_FAILED;
524+ }
525+ return ge::GRAPH_SUCCESS;
526+}
527+ 
528+ge::graphStatus CheckSingleParaDWgate(CompressorGradContext &compressorGradContext)
529+{
530+ if (compressorGradContext.dWgate.desc == nullptr) {
531+ return ge::GRAPH_SUCCESS;
532+ }
533+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(compressorGradContext,
534+ compressorGradContext.dWgate.desc,
535+ D_WGATE_NAME) ||
536+ ge::GRAPH_SUCCESS != CheckDimNumSupport(compressorGradContext.dWgate.shape,
537+ D_WGATE_NAME,
538+ compressorGradContext.opName)) {
539+ return ge::GRAPH_FAILED;
540+ }
541+ return ge::GRAPH_SUCCESS;
542+}
543+ 
544+ge::graphStatus CheckSingleParaDApe(CompressorGradContext &compressorGradContext)
545+{
546+ if (compressorGradContext.dApe.desc == nullptr) {
547+ return ge::GRAPH_SUCCESS;
548+ }
549+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(compressorGradContext,
550+ compressorGradContext.dApe.desc,
551+ D_APE_NAME) ||
552+ ge::GRAPH_SUCCESS != CheckDimNumSupport(compressorGradContext.dApe.shape,
553+ D_APE_NAME,
554+ compressorGradContext.opName)) {
555+ return ge::GRAPH_FAILED;
556+ }
557+ return ge::GRAPH_SUCCESS;
558+}
559+ 
560+ge::graphStatus CheckSingleParaCmpRatio(CompressorGradContext &compressorGradContext)
561+{
562+ if (ge::GRAPH_SUCCESS != CheckAttrValueSupportInterval(compressorGradContext.cmpRatio,
563+ MIN_CMP_RATIO,
564+ MAX_CMP_RATIO,
565+ CMP_RATIO_NAME,
566+ compressorGradContext.opName)) {
567+ return ge::GRAPH_FAILED;
568+ }
569+ return ge::GRAPH_SUCCESS;
570+}
571+ 
572+ge::graphStatus CheckSingleParaCoff(CompressorGradContext &compressorGradContext)
573+{
574+ if (ge::GRAPH_SUCCESS != CheckAttrValueSupportList(compressorGradContext.coff,
575+ COFF,
576+ COFF_NAME,
577+ compressorGradContext.opName)) {
578+ return ge::GRAPH_FAILED;
579+ }
580+ return ge::GRAPH_SUCCESS;
581+}
582+ 
583+ge::graphStatus CheckFeature(CompressorGradContext &compressorGradContext,
584+ CompressorGradBaseParams &compressorGradBaseParams)
585+{
586+ if (ge::GRAPH_SUCCESS != CheckFeatureValueSupport(&compressorGradBaseParams.headDim, HEAD_DIM,
587+ "headDim", compressorGradContext.opName)) {
588+ return ge::GRAPH_FAILED;
589+ }
590+ OP_CHECK_IF(compressorGradBaseParams.hiddenSize > MAX_HIDDEN_SIZE ||
591+ compressorGradBaseParams.hiddenSize < MIN_HIDDEN_SIZE ||
592+ compressorGradBaseParams.hiddenSize % ALIGN_FACTOR_HIDDEN_SIZE != 0,
593+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(compressorGradContext.opName, "hiddenSize",
594+ std::to_string(compressorGradBaseParams.hiddenSize),
595+ "should be within [" + std::to_string(MIN_HIDDEN_SIZE) + ", " +
596+ std::to_string(MAX_HIDDEN_SIZE) + "] and be 512-aligned"),
597+ return ge::GRAPH_FAILED);
598+ return ge::GRAPH_SUCCESS;
599+}
600+ 
601+ge::graphStatus LogErrorShapeConsistency(const std::string &name, const gert::StorageShape *shape,
602+ const uint32_t &dimNum, const std::string &subName,
603+ const uint32_t &expectNum, const char *opName)
604+{
605+ if (shape == nullptr) {
606+ return ge::GRAPH_SUCCESS;
607+ }
608+ 
609+ const uint32_t actualNum = shape->GetStorageShape().GetDim(dimNum);
610+ OP_CHECK_IF(actualNum != expectNum,
611+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(opName, name,
612+ "dim " + std::to_string(dimNum) + "=" + std::to_string(actualNum),
613+ "should be equal to " + subName + ": " + std::to_string(expectNum)),
614+ return ge::GRAPH_FAILED);
615+ 
616+ return ge::GRAPH_SUCCESS;
617+}
618+ 
619+ge::graphStatus CheckShapeConsistency(CompressorGradContext &compressorGradContext,
620+ CompressorGradBaseParams &compressorGradBaseParams)
621+{
622+ uint8_t coff = compressorGradContext.coff == nullptr ? COFF_VALUE :
623+ static_cast<uint8_t>(*compressorGradContext.coff);
624+ auto coffD = compressorGradBaseParams.headDim * coff;
625+ if (ge::GRAPH_SUCCESS !=
626+ LogErrorShapeConsistency("cuSeqlens", compressorGradContext.cuSeqlens.shape, COMPRESSOR_GRAD_DIM_INDEX_0,
627+ "batchSize+1", compressorGradBaseParams.batchSize + 1,
628+ compressorGradContext.opName) ||
629+ ge::GRAPH_SUCCESS !=
630+ LogErrorShapeConsistency("seqUsed", compressorGradContext.seqUsed.shape, COMPRESSOR_GRAD_DIM_INDEX_0,
631+ "batchSize", compressorGradBaseParams.batchSize, compressorGradContext.opName) ||
632+ ge::GRAPH_SUCCESS !=
633+ LogErrorShapeConsistency("startPos", compressorGradContext.startPos.shape, COMPRESSOR_GRAD_DIM_INDEX_0,
634+ "batchSize", compressorGradBaseParams.batchSize, compressorGradContext.opName) ||
635+ ge::GRAPH_SUCCESS !=
636+ LogErrorShapeConsistency("wkv", compressorGradContext.wkv.shape, COMPRESSOR_GRAD_DIM_INDEX_1, "hiddenSize",
637+ compressorGradBaseParams.hiddenSize, compressorGradContext.opName) ||
638+ ge::GRAPH_SUCCESS !=
639+ LogErrorShapeConsistency("wgate", compressorGradContext.wgate.shape, COMPRESSOR_GRAD_DIM_INDEX_1,
640+ "hiddenSize", compressorGradBaseParams.hiddenSize, compressorGradContext.opName) ||
641+ ge::GRAPH_SUCCESS !=
642+ LogErrorShapeConsistency("wkv", compressorGradContext.wkv.shape, COMPRESSOR_GRAD_DIM_INDEX_0,
643+ "coff*headDim", static_cast<uint32_t>(coffD), compressorGradContext.opName) ||
644+ ge::GRAPH_SUCCESS !=
645+ LogErrorShapeConsistency("wgate", compressorGradContext.wgate.shape, COMPRESSOR_GRAD_DIM_INDEX_0,
646+ "coff*headDim", static_cast<uint32_t>(coffD), compressorGradContext.opName)) {
647+ return ge::GRAPH_FAILED;
648+ }
649+ 
650+ if (compressorGradContext.x.shape->GetStorageShape().GetDimNum() == COMPRESSOR_GRAD_DIM_NUM_2 &&
651+ compressorGradContext.dCmpKv.shape->GetStorageShape().GetDimNum() == COMPRESSOR_GRAD_DIM_NUM_2 &&
652+ compressorGradContext.softmaxScore.shape->GetStorageShape().GetDimNum() == COMPRESSOR_GRAD_DIM_NUM_3 &&
653+ compressorGradContext.kv.shape->GetStorageShape().GetDimNum() == COMPRESSOR_GRAD_DIM_NUM_3) {
654+ if (ge::GRAPH_SUCCESS !=
655+ LogErrorShapeConsistency("dCmpKv", compressorGradContext.dCmpKv.shape, COMPRESSOR_GRAD_DIM_INDEX_1,
656+ "headDim", compressorGradBaseParams.headDim, compressorGradContext.opName) ||
657+ ge::GRAPH_SUCCESS !=
658+ LogErrorShapeConsistency("softmaxScore", compressorGradContext.softmaxScore.shape,
659+ COMPRESSOR_GRAD_DIM_INDEX_1, "coff*cmp_ratio",
660+ coff * compressorGradBaseParams.cmpRatio, compressorGradContext.opName) ||
661+ ge::GRAPH_SUCCESS !=
662+ LogErrorShapeConsistency("softmaxScore", compressorGradContext.softmaxScore.shape,
663+ COMPRESSOR_GRAD_DIM_INDEX_2, "headDim",
664+ compressorGradBaseParams.headDim, compressorGradContext.opName) ||
665+ ge::GRAPH_SUCCESS !=
666+ LogErrorShapeConsistency("kv", compressorGradContext.kv.shape,
667+ COMPRESSOR_GRAD_DIM_INDEX_1, "coff*cmp_ratio",
668+ coff * compressorGradBaseParams.cmpRatio, compressorGradContext.opName) ||
669+ ge::GRAPH_SUCCESS !=
670+ LogErrorShapeConsistency("kv", compressorGradContext.kv.shape,
671+ COMPRESSOR_GRAD_DIM_INDEX_2, "headDim",
672+ compressorGradBaseParams.headDim, compressorGradContext.opName)) {
673+ return ge::GRAPH_FAILED;
674+ }
675+ }
676+ 
677+ if (compressorGradContext.x.shape->GetStorageShape().GetDimNum() == COMPRESSOR_GRAD_DIM_NUM_3 &&
678+ compressorGradContext.dCmpKv.shape->GetStorageShape().GetDimNum() == COMPRESSOR_GRAD_DIM_NUM_3 &&
679+ compressorGradContext.softmaxScore.shape->GetStorageShape().GetDimNum() == COMPRESSOR_GRAD_DIM_NUM_4 &&
680+ compressorGradContext.kv.shape->GetStorageShape().GetDimNum() == COMPRESSOR_GRAD_DIM_NUM_4) {
681+ if (ge::GRAPH_SUCCESS !=
682+ LogErrorShapeConsistency("dCmpKv", compressorGradContext.dCmpKv.shape,
683+ COMPRESSOR_GRAD_DIM_INDEX_0, "batchSize",
684+ compressorGradBaseParams.batchSize, compressorGradContext.opName) ||
685+ ge::GRAPH_SUCCESS !=
686+ LogErrorShapeConsistency("dCmpKv", compressorGradContext.dCmpKv.shape,
687+ COMPRESSOR_GRAD_DIM_INDEX_2, "headDim",
688+ compressorGradBaseParams.headDim, compressorGradContext.opName) ||
689+ ge::GRAPH_SUCCESS !=
690+ LogErrorShapeConsistency("softmaxScore", compressorGradContext.softmaxScore.shape,
691+ COMPRESSOR_GRAD_DIM_INDEX_0, "batchSize",
692+ compressorGradBaseParams.batchSize, compressorGradContext.opName) ||
693+ ge::GRAPH_SUCCESS !=
694+ LogErrorShapeConsistency("softmaxScore", compressorGradContext.softmaxScore.shape,
695+ COMPRESSOR_GRAD_DIM_INDEX_2, "coff*cmp_ratio",
696+ coff * compressorGradBaseParams.cmpRatio, compressorGradContext.opName) ||
697+ ge::GRAPH_SUCCESS !=
698+ LogErrorShapeConsistency("softmaxScore", compressorGradContext.softmaxScore.shape,
699+ COMPRESSOR_GRAD_DIM_INDEX_3, "headDim",
700+ compressorGradBaseParams.headDim, compressorGradContext.opName) ||
701+ ge::GRAPH_SUCCESS !=
702+ LogErrorShapeConsistency("kv", compressorGradContext.kv.shape, COMPRESSOR_GRAD_DIM_INDEX_0, "batchSize",
703+ compressorGradBaseParams.batchSize, compressorGradContext.opName) ||
704+ ge::GRAPH_SUCCESS !=
705+ LogErrorShapeConsistency("kv", compressorGradContext.kv.shape,
706+ COMPRESSOR_GRAD_DIM_INDEX_2, "coff*cmp_ratio",
707+ coff * compressorGradBaseParams.cmpRatio, compressorGradContext.opName) ||
708+ ge::GRAPH_SUCCESS !=
709+ LogErrorShapeConsistency("kv", compressorGradContext.kv.shape,
710+ COMPRESSOR_GRAD_DIM_INDEX_3, "headDim",
711+ compressorGradBaseParams.headDim, compressorGradContext.opName)) {
712+ return ge::GRAPH_FAILED;
713+ }
714+ }
715+ return ge::GRAPH_SUCCESS;
716+}
717+ 
718+ge::graphStatus CheckDtypeConsistencyX(const gert::CompileTimeTensorDesc *desc,
719+ const std::string &name, CompressorGradContext &compressorGradContext)
720+{
721+ const auto actualDtype = desc->GetDataType();
722+ OP_CHECK_IF(actualDtype != compressorGradContext.dtype,
723+ OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(compressorGradContext.opName, name,
724+ DataTypeToSerialString(actualDtype),
725+ "should be same with x: " + DataTypeToSerialString(compressorGradContext.dtype)),
726+ return ge::GRAPH_FAILED);
727+ return ge::GRAPH_SUCCESS;
728+}
729+ 
730+ge::graphStatus CheckDtypeConsistency(CompressorGradContext &compressorGradContext)
731+{
732+ if (CheckDtypeConsistencyX(compressorGradContext.wkv.desc,
733+ WKV_NAME,
734+ compressorGradContext) != ge::GRAPH_SUCCESS ||
735+ CheckDtypeConsistencyX(compressorGradContext.wgate.desc,
736+ WGATE_NAME,
737+ compressorGradContext) != ge::GRAPH_SUCCESS ||
738+ CheckDtypeConsistencyX(compressorGradContext.dCmpKv.desc,
739+ D_CMP_KV_NAME,
740+ compressorGradContext) != ge::GRAPH_SUCCESS ||
741+ CheckDtypeConsistencyX(compressorGradContext.dX.desc,
742+ D_X_NAME,
743+ compressorGradContext) != ge::GRAPH_SUCCESS ||
744+ CheckDtypeConsistencyX(compressorGradContext.dWkv.desc,
745+ D_WKV_NAME,
746+ compressorGradContext) != ge::GRAPH_SUCCESS ||
747+ CheckDtypeConsistencyX(compressorGradContext.dWgate.desc,
748+ D_WGATE_NAME,
749+ compressorGradContext) != ge::GRAPH_SUCCESS) {
750+ return ge::GRAPH_FAILED;
751+ }
752+ return ge::GRAPH_SUCCESS;
753+}
754+ 
755+ge::graphStatus CheckDimNumConsistency(CompressorGradContext &compressorGradContext)
756+{
757+ auto xDimNum = compressorGradContext.x.shape->GetStorageShape().GetDimNum();
758+ OP_CHECK_IF(xDimNum != compressorGradContext.dX.shape->GetStorageShape().GetDimNum(),
759+ OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(compressorGradContext.opName, "d_x",
760+ std::to_string(compressorGradContext.dX.shape->GetStorageShape().GetDimNum()),
761+ "dim num should be equal to x: " + std::to_string(xDimNum)),
762+ return ge::GRAPH_FAILED);
763+
764+ OP_CHECK_IF(xDimNum != compressorGradContext.dCmpKv.shape->GetStorageShape().GetDimNum(),
765+ OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(compressorGradContext.opName, "d_cmp_kv",
766+ std::to_string(compressorGradContext.dCmpKv.shape->GetStorageShape().GetDimNum()),
767+ "dim num should be equal to x: " + std::to_string(xDimNum)),
768+ return ge::GRAPH_FAILED);
769+
770+ OP_CHECK_IF(xDimNum != compressorGradContext.softmaxScore.shape->GetStorageShape().GetDimNum() - 1,
771+ OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(compressorGradContext.opName, "softmax_score",
772+ std::to_string(compressorGradContext.softmaxScore.shape->GetStorageShape().GetDimNum()),
773+ "dim num should be x dim + 1: " + std::to_string(xDimNum + 1)),
774+ return ge::GRAPH_FAILED);
775+
776+ OP_CHECK_IF(xDimNum != compressorGradContext.kv.shape->GetStorageShape().GetDimNum() - 1,
777+ OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(compressorGradContext.opName, "kv",
778+ std::to_string(compressorGradContext.kv.shape->GetStorageShape().GetDimNum()),
779+ "dim num should be x dim + 1: " + std::to_string(xDimNum + 1)),
780+ return ge::GRAPH_FAILED);
781+ return ge::GRAPH_SUCCESS;
782+}
783+ 
784+ge::graphStatus CheckBlockDimConstrain(CompressorGradContext &compressorGradContext,
785+ CompressorGradBaseParams &compressorGradBaseParams,
786+ uint32_t &aicNum)
787+{
788+ uint32_t minBlockNum = compressorGradBaseParams.headDim / 128; // 128 is the largest dBaseSize
789+ OP_CHECK_IF(aicNum < minBlockNum,
790+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(compressorGradContext.opName, "aicNum",
791+ std::to_string(aicNum),
792+ "should not be less than " + std::to_string(minBlockNum)),
793+ return ge::GRAPH_FAILED);
794+ return ge::GRAPH_SUCCESS;
795+}
796+ 
797+ge::graphStatus CheckMultiParaConsistency(CompressorGradContext &compressorGradContext,
798+ CompressorGradBaseParams &compressorGradBaseParams)
799+{
800+ if (CheckShapeConsistency(compressorGradContext, compressorGradBaseParams) != ge::GRAPH_SUCCESS ||
801+ CheckDtypeConsistency(compressorGradContext) != ge::GRAPH_SUCCESS ||
802+ CheckDimNumConsistency(compressorGradContext) != ge::GRAPH_SUCCESS) {
803+ return ge::GRAPH_FAILED;
804+ }
805+ return ge::GRAPH_SUCCESS;
806+}
807+ 
808+ge::graphStatus CheckSinglePara(CompressorGradContext &compressorGradContext)
809+{
810+ if (ge::GRAPH_SUCCESS != CheckSingleParaX(compressorGradContext) ||
811+ ge::GRAPH_SUCCESS != CheckSingleParaWkv(compressorGradContext) ||
812+ ge::GRAPH_SUCCESS != CheckSingleParaWgate(compressorGradContext) ||
813+ ge::GRAPH_SUCCESS != CheckSingleParaDCmpKv(compressorGradContext) ||
814+ ge::GRAPH_SUCCESS != CheckSingleParaSoftmaxScore(compressorGradContext) ||
815+ ge::GRAPH_SUCCESS != CheckSingleParaCuSeqlens(compressorGradContext) ||
816+ ge::GRAPH_SUCCESS != CheckSingleParaSeqused(compressorGradContext) ||
817+ ge::GRAPH_SUCCESS != CheckSingleParaStartPos(compressorGradContext) ||
818+ ge::GRAPH_SUCCESS != CheckSingleParaDX(compressorGradContext) ||
819+ ge::GRAPH_SUCCESS != CheckSingleParaDWkv(compressorGradContext) ||
820+ ge::GRAPH_SUCCESS != CheckSingleParaDWgate(compressorGradContext) ||
821+ ge::GRAPH_SUCCESS != CheckSingleParaDApe(compressorGradContext) ||
822+ ge::GRAPH_SUCCESS != CheckSingleParaCmpRatio(compressorGradContext) ||
823+ ge::GRAPH_SUCCESS != CheckSingleParaCoff(compressorGradContext)) {
824+ return ge::GRAPH_FAILED;
825+ }
826+ return ge::GRAPH_SUCCESS;
827+}
828+ 
829+ge::graphStatus CheckRequiredInOutExistence(CompressorGradContext &context)
830+{
831+ OP_CHECK_IF(context.x.shape == nullptr,
832+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "x", "shape is nullptr"),
833+ return ge::GRAPH_FAILED);
834+ OP_CHECK_IF(context.x.desc == nullptr,
835+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "x", "desc is nullptr"),
836+ return ge::GRAPH_FAILED);
837+ OP_CHECK_IF(context.wkv.shape == nullptr,
838+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "wkv", "shape is nullptr"),
839+ return ge::GRAPH_FAILED);
840+ OP_CHECK_IF(context.wkv.desc == nullptr,
841+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "wkv", "desc is nullptr"),
842+ return ge::GRAPH_FAILED);
843+ OP_CHECK_IF(context.wgate.shape == nullptr,
844+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "wgate", "shape is nullptr"),
845+ return ge::GRAPH_FAILED);
846+ OP_CHECK_IF(context.wgate.desc == nullptr,
847+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "wgate", "desc is nullptr"),
848+ return ge::GRAPH_FAILED);
849+ OP_CHECK_IF(context.dCmpKv.shape == nullptr,
850+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "d_cmp_kv", "shape is nullptr"),
851+ return ge::GRAPH_FAILED);
852+ OP_CHECK_IF(context.dCmpKv.desc == nullptr,
853+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "d_cmp_kv", "desc is nullptr"),
854+ return ge::GRAPH_FAILED);
855+ OP_CHECK_IF(context.softmaxScore.shape == nullptr,
856+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "softmax_score", "shape is nullptr"),
857+ return ge::GRAPH_FAILED);
858+ OP_CHECK_IF(context.softmaxScore.desc == nullptr,
859+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "softmax_score", "desc is nullptr"),
860+ return ge::GRAPH_FAILED);
861+ OP_CHECK_IF(context.kv.shape == nullptr,
862+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "kv", "shape is nullptr"),
863+ return ge::GRAPH_FAILED);
864+ OP_CHECK_IF(context.kv.desc == nullptr,
865+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "kv", "desc is nullptr"),
866+ return ge::GRAPH_FAILED);
867+ 
868+ OP_CHECK_IF(context.dX.shape == nullptr,
869+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "d_x", "shape is nullptr"),
870+ return ge::GRAPH_FAILED);
871+ OP_CHECK_IF(context.dX.desc == nullptr,
872+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "d_x", "desc is nullptr"),
873+ return ge::GRAPH_FAILED);
874+
875+ OP_CHECK_IF(context.dWkv.shape == nullptr,
876+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "d_wkv", "shape is nullptr"),
877+ return ge::GRAPH_FAILED);
878+ OP_CHECK_IF(context.dWkv.desc == nullptr,
879+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "d_wkv", "desc is nullptr"),
880+ return ge::GRAPH_FAILED);
881+ 
882+ OP_CHECK_IF(context.dWgate.shape == nullptr,
883+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "d_wgate", "shape is nullptr"),
884+ return ge::GRAPH_FAILED);
885+ OP_CHECK_IF(context.dWgate.desc == nullptr,
886+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "d_wgate", "desc is nullptr"),
887+ return ge::GRAPH_FAILED);
888+ 
889+ OP_CHECK_IF(context.dApe.shape == nullptr,
890+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "d_ape", "shape is nullptr"),
891+ return ge::GRAPH_FAILED);
892+ OP_CHECK_IF(context.dApe.desc == nullptr,
893+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "d_ape", "desc is nullptr"),
894+ return ge::GRAPH_FAILED);
895+
896+ if (context.layout == LayoutType::LAYOUT_TH) {
897+ OP_CHECK_IF(context.cuSeqlens.desc == nullptr,
898+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "cu_seqlens",
899+ "in TH layout, should not be nullptr"),
900+ return ge::GRAPH_FAILED);
901+ OP_CHECK_IF(context.cuSeqlens.shape == nullptr,
902+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "cu_seqlens",
903+ "in TH layout, should not be nullptr"),
904+ return ge::GRAPH_FAILED);
905+ } else {
906+ OP_CHECK_IF(context.cuSeqlens.desc != nullptr,
907+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "cu_seqlens",
908+ "in BSH layout, must be nullptr"),
909+ return ge::GRAPH_FAILED);
910+ OP_CHECK_IF(context.cuSeqlens.shape != nullptr,
911+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "cu_seqlens",
912+ "in BSH layout, must be nullptr"),
913+ return ge::GRAPH_FAILED);
914+ }
915+ return ge::GRAPH_SUCCESS;
916+}
917+ 
918+ge::graphStatus CheckRequiredAttrExistence(CompressorGradContext &context)
919+{
920+ OP_CHECK_IF(context.cmpRatio == nullptr,
921+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.opName, "cmp_ratio", "attr is nullptr"),
922+ return ge::GRAPH_FAILED);
923+ return ge::GRAPH_SUCCESS;
924+}
925+ 
926+ge::graphStatus CheckRequiredParaExistence(CompressorGradContext &compressorGradcontext)
927+{
928+ if (CheckRequiredInOutExistence(compressorGradcontext) != ge::GRAPH_SUCCESS ||
929+ CheckRequiredAttrExistence(compressorGradcontext) != ge::GRAPH_SUCCESS) {
930+ return ge::GRAPH_FAILED;
931+ }
932+ return ge::GRAPH_SUCCESS;
933+}
934+ 
935+ge::graphStatus CheckEmptyTensor(CompressorGradContext &compressorGradcontext)
936+{
937+ // CompressorGrad 不支持空 tensor：与正向不同（正向 x 支持 B/S/T=0 走 EMPTY_X 分支），
938+ // 反向无空 tensor 分支——所有输入/输出 shapeSize 必须 > 0，空则直接拦截
939+ if (compressorGradcontext.x.shape->GetStorageShape().GetShapeSize() == 0 ||
940+ compressorGradcontext.wkv.shape->GetStorageShape().GetShapeSize() == 0 ||
941+ compressorGradcontext.wgate.shape->GetStorageShape().GetShapeSize() == 0 ||
942+ compressorGradcontext.dCmpKv.shape->GetStorageShape().GetShapeSize() == 0 ||
943+ compressorGradcontext.softmaxScore.shape->GetStorageShape().GetShapeSize() == 0 ||
944+ compressorGradcontext.kv.shape->GetStorageShape().GetShapeSize() == 0 ||
945+ compressorGradcontext.dX.shape->GetStorageShape().GetShapeSize() == 0 ||
946+ compressorGradcontext.dWkv.shape->GetStorageShape().GetShapeSize() == 0 ||
947+ compressorGradcontext.dWgate.shape->GetStorageShape().GetShapeSize() == 0 ||
948+ compressorGradcontext.dApe.shape->GetStorageShape().GetShapeSize() == 0) {
949+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(compressorGradcontext.opName, "x", "0",
950+ "CompressorGrad does not support empty tensor: all inputs/outputs shapeSize must be > 0");
951+ return ge::GRAPH_FAILED;
952+ }
953+ 
954+ return ge::GRAPH_SUCCESS;
955+}
956+ 
957+static ge::graphStatus CompressorGradTilingFunc(gert::TilingContext *context)
958+{
959+ CompressorGradContext compressorGradContext{};
960+ CompressorGradBaseParams compressorGradBaseParams{};
961+ CompressorGradWorkspaceParams compressorGradWorkspaceParams{};
962+ OP_CHECK_IF(context == nullptr,
963+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON("CompressorGrad", "context", "is nullptr"),
964+ return ge::GRAPH_FAILED);
965+ 
966+ OP_LOGI("Getting Tiling");
967+ 
968+ OP_CHECK_IF(context->GetPlatformInfo() == nullptr,
969+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context->GetNodeName(), "platformInfo", "is nullptr"),
970+ return ge::GRAPH_FAILED);
971+ 
972+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
973+ 
974+ uint32_t aivNum = ascendcPlatform.GetCoreNumAiv();
975+ uint32_t aicNum = ascendcPlatform.GetCoreNumAic();
976+ OP_CHECK_IF(aicNum == 0 || aivNum == 0,
977+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context->GetNodeName(), "aicNum/aivNum",
978+ "num of core obtained is 0"),
979+ return ge::GRAPH_FAILED);
980+
981+ uint32_t blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum);
982+ context->SetBlockDim(blockDim);
983+
984+ if (ConvertContext(*context, compressorGradContext) != ge::GRAPH_SUCCESS) {
985+ OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context->GetNodeName(), "context",
986+ "error occurred while converting tilingContext to CompressorGrad context");
987+ return ge::GRAPH_FAILED;
988+ }
989+ 
990+ compressorGradContext.blockDim = aicNum;
991+ if (CheckRequiredParaExistence(compressorGradContext) != ge::GRAPH_SUCCESS) {
992+ return ge::GRAPH_FAILED;
993+ }
994+ 
995+ if (CheckEmptyTensor(compressorGradContext) != ge::GRAPH_SUCCESS) {
996+ return ge::GRAPH_FAILED;
997+ }
998+ 
999+ if (CheckSinglePara(compressorGradContext) != ge::GRAPH_SUCCESS) {
1000+ return ge::GRAPH_FAILED;
1001+ }
1002+ 
1003+ if (SetBaseInfo(compressorGradContext, compressorGradBaseParams, aicNum) != ge::GRAPH_SUCCESS) {
1004+ return ge::GRAPH_FAILED;
1005+ }
1006+ 
1007+ if (CheckFeature(compressorGradContext, compressorGradBaseParams) != ge::GRAPH_SUCCESS) {
1008+ return ge::GRAPH_FAILED;
1009+ }
1010+ 
1011+ if (CheckMultiParaConsistency(compressorGradContext, compressorGradBaseParams) != ge::GRAPH_SUCCESS) {
1012+ return ge::GRAPH_FAILED;
1013+ }
1014+ 
1015+ if (CheckBlockDimConstrain(compressorGradContext, compressorGradBaseParams, aicNum) != ge::GRAPH_SUCCESS) {
1016+ return ge::GRAPH_FAILED;
1017+ }
1018+ 
1019+ size_t libapiSize = 0;
1020+ size_t workspaceSize = 0;
1021+ 
1022+ if (CalcWorkSpace(compressorGradContext, compressorGradBaseParams, compressorGradWorkspaceParams, workspaceSize,
1023+ libapiSize, aicNum) != ge::GRAPH_SUCCESS) {
1024+ return ge::GRAPH_FAILED;
1025+ }
1026+ if (GenTilingKey(compressorGradContext) != ge::GRAPH_SUCCESS) {
1027+ return ge::GRAPH_FAILED;
1028+ }
1029+ 
1030+ uint8_t coff = compressorGradContext.coff == nullptr ? COFF_VALUE :
1031+ static_cast<uint8_t>(*compressorGradContext.coff);
1032+ CompressorGradTiling *tilingData = context->GetTilingData<CompressorGradTiling>();
1033+ tilingData->batch_size = compressorGradBaseParams.batchSize;
1034+ tilingData->token_size = compressorGradBaseParams.tokenSize;
1035+ tilingData->seq_size = compressorGradBaseParams.seqSize;
1036+ tilingData->cmp_ratio = compressorGradBaseParams.cmpRatio;
1037+ tilingData->hidden_size = compressorGradBaseParams.hiddenSize;
1038+ tilingData->head_dim = compressorGradBaseParams.headDim;
1039+ // ── 核数（与 launch block_dim 一致；vec 每核 2 子核）──
1040+ tilingData->cube_core_num = compressorGradContext.blockDim;
1041+ tilingData->core_num = compressorGradContext.blockDim * 2;
1042+ // ── shape 派生 ──
1043+ tilingData->total_head_dim = coff * compressorGradBaseParams.headDim;
1044+ tilingData->cmp_row_cnt = coff * compressorGradBaseParams.cmpRatio;
1045+ tilingData->cmp_size = coff * compressorGradBaseParams.cmpRatio * compressorGradBaseParams.headDim;
1046+ tilingData->cmp_kv_batch_stride =
1047+ (compressorGradBaseParams.seqSize + compressorGradBaseParams.cmpRatio - 1) / compressorGradBaseParams.cmpRatio;
1048+ if (compressorGradContext.layout == LayoutType::LAYOUT_BSH) {
1049+ tilingData->x_rows = compressorGradBaseParams.batchSize * compressorGradBaseParams.seqSize;
1050+ tilingData->cmp_kv_rows = compressorGradBaseParams.batchSize * tilingData->cmp_kv_batch_stride;
1051+ } else {
1052+ tilingData->x_rows = compressorGradBaseParams.tokenSize;
1053+ tilingData->cmp_kv_rows = std::min(compressorGradBaseParams.tokenSize,
1054+ compressorGradBaseParams.tokenSize /
1055+ compressorGradBaseParams.cmpRatio +compressorGradBaseParams.batchSize);
1056+ }
1057+ // ── 分核派生 ──
1058+ tilingData->group_size = compressorGradBaseParams.headDim / 128;
1059+ tilingData->group_num = compressorGradContext.blockDim / tilingData->group_size;
1060+ tilingData->cube_m_base_size = 128 * (2 / coff);
1061+ // coff=1 时每 group 每轮只布置 2*dealScNum 块（保证子槽 dealScNum*cmpRatio <= 128，不超 L1/L0 物理行）
1062+ tilingData->deal_sc_num = 128 / compressorGradBaseParams.cmpRatio;
1063+ tilingData->group_deal_sc_num = tilingData->deal_sc_num * (2 / coff);
1064+ tilingData->total_sc_num_per_round = tilingData->group_num * tilingData->group_deal_sc_num;
1065+ // xArrangeGm 每 group 实际行数 = 数据行(gs 块 × cr) + coff=2 时 1 个 cr 头部（紧凑布局）
1066+ tilingData->group_row_stride = tilingData->group_deal_sc_num * compressorGradBaseParams.cmpRatio +
1067+ (coff - 1) * compressorGradBaseParams.cmpRatio;
1068+ tilingData->db_row_cnt = tilingData->group_num * tilingData->group_row_stride;
1069+ // ── 编译期派生（TilingKey 折叠值，与 kernel 内联算术恒等）──
1070+ tilingData->coff_coef = 2 / coff;
1071+ tilingData->d_deal_size = 128 / coff;
1072+ tilingData->m_deal_size = 128 * coff;
1073+ // ── workspace 分区（FP32 元素数；ape/dW 单缓冲，dX/x/dXCache 双缓冲 dbRatio=2）──
1074+ tilingData->dape_ws_size = tilingData->group_num * tilingData->cmp_size * tilingData->coff_coef;
1075+ tilingData->d_x_ws_size = 2 * tilingData->cube_core_num * 256 * compressorGradBaseParams.hiddenSize;
1076+ tilingData->d_w_weight_ws_size = tilingData->group_num * tilingData->total_head_dim *
1077+ compressorGradBaseParams.hiddenSize;
1078+ tilingData->x_ws_size = 2 * tilingData->group_num * tilingData->group_row_stride *
1079+ compressorGradBaseParams.hiddenSize;
1080+ tilingData->d_x_cache_ws_size = 2 * compressorGradBaseParams.cmpRatio * compressorGradBaseParams.hiddenSize;
1081+ 
1082+ context->SetTilingKey(compressorGradContext.tilingKey);
1083+ context->SetBlockDim(compressorGradContext.blockDim);
1084+ OP_LOGI(compressorGradContext.opName, "block dim: %u.", compressorGradContext.blockDim);
1085+ return ge::GRAPH_SUCCESS;
1086+}
1087+ 
1088+static ge::graphStatus TilingParseForCompressorGrad(gert::TilingParseContext *context)
1089+{
1090+ (void)context;
1091+ return ge::GRAPH_SUCCESS;
1092+}
1093+ 
1094+IMPL_OP_OPTILING(CompressorGrad)
1095+ .Tiling(CompressorGradTilingFunc)
1096+ .TilingParse<CompressorGradCompileInfo>(TilingParseForCompressorGrad);
1097+} // namespace optiling
Aattention/compressor_grad/op_host/compressor_grad_tiling.h+304-0
@@ -0,0 +1,304 @@
1+/* *
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+  */
10+ 
11+/*!
12+ * \file compressor_grad_tiling.h
13+ * \brief
14+ */
15+ 
16+#ifndef COMPRESSOR_GRAD_TILING_H_
17+#define COMPRESSOR_GRAD_TILING_H_
18+ 
19+#include "exe_graph/runtime/tiling_context.h"
20+#include "tiling/platform/platform_ascendc.h"
21+#include "register/op_def_registry.h"
22+#include "register/tilingdata_base.h"
23+#include "tiling/tiling_api.h"
24+ 
25+namespace optiling {
26+// INPUT
27+constexpr uint32_t TOKEN_X_INPUT_INDEX = 0;
28+constexpr uint32_t WEIGHT_KV_INPUT_INDEX = 1;
29+constexpr uint32_t WEIGHT_WGATE_INPUT_INDEX = 2;
30+ 
31+constexpr uint32_t D_CMP_KV_INPUT_INDEX = 3;
32+ 
33+constexpr uint32_t SOFTMAX_SCORE_INPUT_INDEX = 4;
34+ 
35+// INPUT(OPTION)
36+constexpr uint32_t KV_INPUT_INDEX = 5;
37+ 
38+constexpr uint32_t CU_SEQ_LEN_INPUT_INDEX = 6;
39+constexpr uint32_t SEQ_USED_INPUT_INDEX = 7;
40+constexpr uint32_t START_POS_INPUT_INDEX = 8;
41+ 
42+// ATTR
43+constexpr uint32_t CMP_RATIO_ATTR_INDEX = 0;
44+constexpr uint32_t COFF_ATTR_INDEX = 1;
45+ 
46+// OUTPUT
47+constexpr uint32_t D_X_OUTPUT_INDEX = 0;
48+constexpr uint32_t D_WKV_OUTPUT_INDEX = 1;
49+constexpr uint32_t D_WGATE_OUTPUT_INDEX = 2;
50+constexpr uint32_t D_APE_OUTPUT_INDEX = 3;
51+ 
52+// ATTR DEFAULT VALUE
53+constexpr uint32_t CMP_RATIO_VALUE = 4;
54+constexpr uint32_t COFF_VALUE = 1;
55+ 
56+constexpr uint32_t COMPRESSOR_GRAD_DIM_NUM_1 = 1;
57+constexpr uint32_t COMPRESSOR_GRAD_DIM_NUM_2 = 2;
58+constexpr uint32_t COMPRESSOR_GRAD_DIM_NUM_3 = 3;
59+constexpr uint32_t COMPRESSOR_GRAD_DIM_NUM_4 = 4;
60+constexpr uint32_t COMPRESSOR_GRAD_DIM_INDEX_0 = 0;
61+constexpr uint32_t COMPRESSOR_GRAD_DIM_INDEX_1 = 1;
62+constexpr uint32_t COMPRESSOR_GRAD_DIM_INDEX_2 = 2;
63+constexpr uint32_t COMPRESSOR_GRAD_DIM_INDEX_3 = 3;
64+ 
65+// CONSTRAINTS
66+constexpr uint32_t MAX_HIDDEN_SIZE = 10240;
67+constexpr uint32_t MIN_HIDDEN_SIZE = 1024;
68+constexpr uint32_t ALIGN_FACTOR_HIDDEN_SIZE = 512;
69+constexpr uint32_t MIN_CMP_RATIO = 2;
70+constexpr uint32_t MAX_CMP_RATIO = 128;
71+ 
72+constexpr uint32_t BATCH_MODE_SCHEDULE = 1;
73+const uint32_t CMP_MAX_AIC_CORE_NUM = 36;
74+ 
75+static const std::string X_NAME = "query";
76+static const std::string WKV_NAME = "wkv";
77+static const std::string WGATE_NAME = "wgate";
78+static const std::string D_CMP_KV_NAME = "d_cmp_kv";
79+static const std::string SOFTMAX_SCORE_NAME = "softmax_score";
80+static const std::string KV_NAME = "kv";
81+static const std::string CU_SEQLENS_NAME = "cu_seqlens";
82+static const std::string SEQUSED_NAME = "seq_used";
83+static const std::string START_POS_NAME = "start_pos";
84+static const std::string CMP_RATIO_NAME = "cmp_ratio";
85+static const std::string COFF_NAME = "coff";
86+static const std::string D_X_NAME = "d_x";
87+static const std::string D_WKV_NAME = "d_wkv";
88+static const std::string D_WGATE_NAME = "d_wgate";
89+static const std::string D_APE_NAME = "d_ape";
90+ 
91+static std::string DataTypeToSerialString(ge::DataType type);
92+ 
93+const std::map<std::string, std::vector<ge::DataType>> DTYPE_SUPPORT_MAP = {
94+ {X_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
95+ {WKV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
96+ {WGATE_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
97+ {D_CMP_KV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
98+ {SOFTMAX_SCORE_NAME, {ge::DT_FLOAT}},
99+ {KV_NAME, {ge::DT_FLOAT}},
100+ {CU_SEQLENS_NAME, {ge::DT_INT32}},
101+ {SEQUSED_NAME, {ge::DT_INT32}},
102+ {START_POS_NAME, {ge::DT_INT32}},
103+ {D_X_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
104+ {D_WKV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
105+ {D_WGATE_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
106+ {D_APE_NAME, {ge::DT_FLOAT}}
107+};
108+ 
109+const std::map<std::string, std::vector<uint32_t>> DIM_NUM_MAP = {
110+ {X_NAME, {COMPRESSOR_GRAD_DIM_NUM_2, COMPRESSOR_GRAD_DIM_NUM_3}},
111+ {WKV_NAME, {COMPRESSOR_GRAD_DIM_NUM_2}},
112+ {WGATE_NAME, {COMPRESSOR_GRAD_DIM_NUM_2}},
113+ {D_CMP_KV_NAME, {COMPRESSOR_GRAD_DIM_NUM_2, COMPRESSOR_GRAD_DIM_NUM_3}},
114+ {SOFTMAX_SCORE_NAME, {COMPRESSOR_GRAD_DIM_NUM_3, COMPRESSOR_GRAD_DIM_NUM_4}},
115+ {KV_NAME, {COMPRESSOR_GRAD_DIM_NUM_3, COMPRESSOR_GRAD_DIM_NUM_4}},
116+ {CU_SEQLENS_NAME, {COMPRESSOR_GRAD_DIM_NUM_1}},
117+ {SEQUSED_NAME, {COMPRESSOR_GRAD_DIM_NUM_1}},
118+ {START_POS_NAME, {COMPRESSOR_GRAD_DIM_NUM_1}},
119+ {D_X_NAME, {COMPRESSOR_GRAD_DIM_NUM_2, COMPRESSOR_GRAD_DIM_NUM_3}},
120+ {D_WKV_NAME, {COMPRESSOR_GRAD_DIM_NUM_2}},
121+ {D_WGATE_NAME, {COMPRESSOR_GRAD_DIM_NUM_2}},
122+ {D_APE_NAME, {COMPRESSOR_GRAD_DIM_NUM_2}}
123+};
124+ 
125+static const std::map<std::string, uint32_t> LAYOUT_DIM_MAP = {
126+ {"BSH", COMPRESSOR_GRAD_DIM_NUM_3},
127+ {"TH", COMPRESSOR_GRAD_DIM_NUM_2},
128+};
129+ 
130+const std::map<ge::DataType, std::string> DATATYPE_TO_STRING_MAP = {
131+ {ge::DT_UNDEFINED, "DT_UNDEFINED"}, // Used to indicate a DataType field has not been set.
132+ {ge::DT_FLOAT, "DT_FLOAT"}, // float type
133+ {ge::DT_FLOAT16, "DT_FLOAT16"}, // fp16 type
134+ {ge::DT_INT8, "DT_INT8"}, // int8 type
135+ {ge::DT_INT16, "DT_INT16"}, // int16 type
136+ {ge::DT_UINT16, "DT_UINT16"}, // uint16 type
137+ {ge::DT_UINT8, "DT_UINT8"}, // uint8 type
138+ {ge::DT_INT32, "DT_INT32"}, // uint32 type
139+ {ge::DT_INT64, "DT_INT64"}, // int64 type
140+ {ge::DT_UINT32, "DT_UINT32"}, // unsigned int32
141+ {ge::DT_UINT64, "DT_UINT64"}, // unsigned int64
142+ {ge::DT_BOOL, "DT_BOOL"}, // bool type
143+ {ge::DT_DOUBLE, "DT_DOUBLE"}, // double type
144+ {ge::DT_DUAL, "DT_DUAL"}, // dual output type
145+ {ge::DT_DUAL_SUB_INT8, "DT_DUAL_SUB_INT8"}, // dual output int8 type
146+ {ge::DT_DUAL_SUB_UINT8, "DT_DUAL_SUB_UINT8"}, // dual output uint8 type
147+ {ge::DT_COMPLEX32, "DT_COMPLEX32"}, // complex32 type
148+ {ge::DT_COMPLEX64, "DT_COMPLEX64"}, // complex64 type
149+ {ge::DT_COMPLEX128, "DT_COMPLEX128"}, // complex128 type
150+ {ge::DT_QINT8, "DT_QINT8"}, // qint8 type
151+ {ge::DT_QINT16, "DT_QINT16"}, // qint16 type
152+ {ge::DT_QINT32, "DT_QINT32"}, // qint32 type
153+ {ge::DT_QUINT8, "DT_QUINT8"}, // quint8 type
154+ {ge::DT_QUINT16, "DT_QUINT16"}, // quint16 type
155+ {ge::DT_RESOURCE, "DT_RESOURCE"}, // resource type
156+ {ge::DT_STRING_REF, "DT_STRING_REF"}, // string ref type
157+ {ge::DT_STRING, "DT_STRING"}, // string type
158+ {ge::DT_VARIANT, "DT_VARIANT"}, // dt_variant type
159+ {ge::DT_BF16, "DT_BFLOAT16"}, // dt_bfloat16 type
160+ {ge::DT_INT4, "DT_INT4"}, // dt_variant type
161+ {ge::DT_UINT1, "DT_UINT1"}, // dt_variant type
162+ {ge::DT_INT2, "DT_INT2"}, // dt_variant type
163+ {ge::DT_UINT2, "DT_UINT2"} // dt_variant type
164+};
165+ 
166+struct RequiredParaInfo {
167+ const gert::CompileTimeTensorDesc *desc;
168+ const gert::StorageShape *shape;
169+};
170+ 
171+struct OptionalParaInfo {
172+ const gert::CompileTimeTensorDesc *desc;
173+ const gert::StorageShape *shape;
174+ const gert::Tensor *tensor;
175+};
176+ 
177+const std::vector<uint32_t> COFF {1, 2};
178+const std::vector<uint32_t> HEAD_DIM {128, 512};
179+ 
180+enum class LayoutType {
181+ LAYOUT_BSH,
182+ LAYOUT_TH
183+};
184+ 
185+enum class TemplateId:uint8_t {
186+ NORMAL = 0,
187+ EMPTY_X = 1,
188+ FULL_LOAD = 2
189+};
190+ 
191+struct CompressorGradContext {
192+ const char *opName;
193+ const char *opType;
194+ fe::PlatFormInfos *platformInfo;
195+
196+ RequiredParaInfo x;
197+ RequiredParaInfo wkv;
198+ RequiredParaInfo wgate;
199+ RequiredParaInfo dCmpKv;
200+ RequiredParaInfo softmaxScore;
201+ OptionalParaInfo kv;
202+ OptionalParaInfo cuSeqlens;
203+ OptionalParaInfo seqUsed;
204+ OptionalParaInfo startPos;
205+ RequiredParaInfo dX;
206+ RequiredParaInfo dWkv;
207+ RequiredParaInfo dWgate;
208+ RequiredParaInfo dApe;
209+ 
210+ const uint32_t *coff;
211+ const uint32_t *cmpRatio;
212+ TemplateId templateId;
213+ 
214+ ge::DataType dtype = ge::DT_BF16;
215+ LayoutType layout = LayoutType::LAYOUT_BSH;
216+ 
217+ size_t *workSpaces;
218+ uint64_t tilingKey;
219+ uint32_t blockDim;
220+};
221+ 
222+struct CompressorGradSplitCoreParams {
223+ uint32_t mStart;
224+ uint32_t mEnd;
225+ uint32_t nStart;
226+ uint32_t nEnd;
227+ uint32_t kStart;
228+ uint32_t kEnd;
229+ };
230+ 
231+// 1. 基础参数结构体
232+struct CompressorGradBaseParams {
233+ uint32_t batchSize = 0; // bastch size（批大小）
234+ uint32_t seqSize = 0; // sequence size（kvs大小）
235+ uint32_t hiddenSize = 0; // hidden size（隐藏层大小）
236+ uint32_t tokenSize = 0; // token size = batchSize * seqSize(token总数：批大小x序列1长度)
237+ uint32_t headDim = 0; // head size of kv
238+ uint32_t featureDim = 0; // head size of kv
239+ uint32_t csSize = 0; // Compress sequence len
240+ uint32_t cmpRatio = 4; // Compress ratio
241+ uint32_t usedCoreNum = 0; // 使用核数
242+ uint32_t nSize = 0; // 预留字段（当前未参与 tiling 决策）
243+ uint64_t stateCacheStrideDim0 = 0; // stateCache第0维的stride
244+ uint32_t kBaseNum = 0;
245+ uint32_t kBaseSize = 0;
246+ uint32_t coreGroupNum = 0;
247+ uint32_t mLoopNum = 0;
248+ CompressorGradSplitCoreParams splitCoreParam[CMP_MAX_AIC_CORE_NUM];
249+};
250+ 
251+struct CompressorGradInnerSplitParams {
252+ uint32_t mBaseSize;
253+ uint32_t dBaseSize;
254+};
255+ 
256+struct CompressorGradWorkspaceParams {
257+ uint32_t mm1KvResSize;
258+ uint32_t mm1ScoreResSize;
259+ uint32_t vec1ResSize;
260+ uint32_t vec1TailCacheSize;
261+ uint32_t dbWorkspaceRatio = 1;
262+};
263+ 
264+BEGIN_TILING_DATA_DEF(CompressorGradTilingData)
265+TILING_DATA_FIELD_DEF(int64_t, batch_size)
266+TILING_DATA_FIELD_DEF(int64_t, token_size)
267+TILING_DATA_FIELD_DEF(int64_t, seq_size)
268+TILING_DATA_FIELD_DEF(int64_t, cmp_ratio)
269+TILING_DATA_FIELD_DEF(int64_t, hidden_size)
270+TILING_DATA_FIELD_DEF(int64_t, head_dim)
271+TILING_DATA_FIELD_DEF(int64_t, cube_core_num)
272+TILING_DATA_FIELD_DEF(int64_t, core_num)
273+TILING_DATA_FIELD_DEF(int64_t, total_head_dim)
274+TILING_DATA_FIELD_DEF(int64_t, cmp_row_cnt)
275+TILING_DATA_FIELD_DEF(int64_t, cmp_size)
276+TILING_DATA_FIELD_DEF(int64_t, cmp_kv_batch_stride)
277+TILING_DATA_FIELD_DEF(int64_t, cmp_kv_rows)
278+TILING_DATA_FIELD_DEF(int64_t, x_rows)
279+TILING_DATA_FIELD_DEF(int64_t, group_size)
280+TILING_DATA_FIELD_DEF(int64_t, group_num)
281+TILING_DATA_FIELD_DEF(int64_t, group_deal_sc_num)
282+TILING_DATA_FIELD_DEF(int64_t, deal_sc_num)
283+TILING_DATA_FIELD_DEF(int64_t, total_sc_num_per_round)
284+TILING_DATA_FIELD_DEF(int64_t, db_row_cnt)
285+TILING_DATA_FIELD_DEF(int64_t, group_row_stride)
286+TILING_DATA_FIELD_DEF(int64_t, coff_coef)
287+TILING_DATA_FIELD_DEF(int64_t, cube_m_base_size)
288+TILING_DATA_FIELD_DEF(int64_t, d_deal_size)
289+TILING_DATA_FIELD_DEF(int64_t, m_deal_size)
290+TILING_DATA_FIELD_DEF(int64_t, dape_ws_size)
291+TILING_DATA_FIELD_DEF(int64_t, dx_ws_size)
292+TILING_DATA_FIELD_DEF(int64_t, d_weight_ws_size)
293+TILING_DATA_FIELD_DEF(int64_t, x_ws_size)
294+TILING_DATA_FIELD_DEF(int64_t, dx_cache_ws_size)
295+END_TILING_DATA_DEF
296+REGISTER_TILING_DATA_CLASS(CompressorGrad, CompressorGradTilingData)
297+ 
298+struct CompressorGradCompileInfo {
299+ int64_t core_num;
300+};
301+ 
302+} // namespace optiling
303+ 
304+#endif // COMPRESSOR_GRAD_TILING_H_
Aattention/compressor_grad/op_kernel/compressor_grad.py+1895-0
Aattention/compressor_grad/tests/CMakeLists.txt+16-0
@@ -0,0 +1,16 @@
1+# -----------------------------------------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# -----------------------------------------------------------------------------------------------------------
10+ 
11+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
12+foreach(SUB_DIR ${CURRENT_DIRS})
13+ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
14+ add_subdirectory(${SUB_DIR})
15+ endif()
16+endforeach()
Aattention/compressor_grad/tests/pytest/compressor_forward.py+449-0
@@ -0,0 +1,449 @@
1+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
2+"""
3+Compressor 算子核心实现：前向 + 反向（纯 PyTorch autograd）。
4+ 
5+计算流程（参考 compressor_golden.py 中的 cpu_compressor）：
6+ 1. new_kv = x @ wkv.T # [*, coff * head_dim]
7+ 2. new_score = x @ wgate.T # [*, coff * head_dim]
8+ 3. new_score += ape（按位置，位置索引 = 全局序列位置 % cmp_ratio）
9+ 4. 写入 page cache（kv_state / score_state，score 已带 ape）
10+ 5. 对每个可压缩块（cmp_ratio 个 token）：
11+ 收集 coff * cmp_ratio 个 token（当前块 + 可选的前一块）
12+ 对 score 按列做 softmax → 权重
13+ compressed = sum(kv * weight, dim=0) → (1, head_dim)
14+ 6. 返回堆叠后的 compressed KV
15+ 
16+反向传播链：
17+ - 通过逐元素乘法和求和
18+ - 通过按列 softmax
19+ - 通过 page-cache 读取 → new_kv / new_score / 旧 state
20+ - 通过 matmul → x, wkv, wgate
21+ - 通过加法 → ape
22+ 
23+精度标准（来自 compressor_golden.py 的 check_result）：
24+ bfloat16: rtol=0.0078125, atol=0.0001, 通过率≥99.5%
25+ 默认: rtol=0.005, atol=0.000025, 通过率≥99.5%
26+"""
27+ 
28+import torch
29+from typing import Optional, List, Tuple
30+ 
31+ 
32+# ===========================================================================
33+# 工具函数
34+# ===========================================================================
35+ 
36+def _softmax_by_column(z: torch.Tensor) -> torch.Tensor:
37+ """
38+ 按列计算 softmax（沿 dim=0，每列独立）。
39+ 输入: (N, D)
40+ 输出: (N, D)，每列之和为 1。
41+ """
42+ z_max = z.max(dim=0, keepdim=True).values
43+ exp_z = torch.exp(z - z_max)
44+ return exp_z / exp_z.sum(dim=0, keepdim=True)
45+ 
46+ 
47+class _QuantGrad(torch.autograd.Function):
S
Sshasha_an16 天前

【一般】python代码风格应该是什么，有的小写加下划线，有的驼峰风格

likedislike
莫允扬
莫允扬
15 天前 评论：
48+ """前向恒等、反向梯度量化：模拟 kernel 反向中间量（dkv/dsb）经
49+ FP16/BF16 存储后回 FP32（与 compressor_grad_cpu_golden 的
50+ d_new_kv.to(iodtype).float() 对齐）。"""
51+ 
52+ @staticmethod
53+ def forward(ctx, x, quant_dtype):
54+ ctx.quant_dtype = quant_dtype
55+ return x
56+ 
57+ @staticmethod
58+ def backward(ctx, grad_output):
59+ return grad_output.to(ctx.quant_dtype).float(), None
60+ 
61+ 
62+def _read_cache(
63+ state: torch.Tensor, # (block_num, block_size, coff * head_dim)
64+ block_table: torch.Tensor, # (B, max_blocks) int32
65+ b_idx: int, # 当前 batch 索引
66+ start_seq: int, # 起始序列位置
67+ end_seq: int, # 结束序列位置
68+ d_start: int, # 特征维度起始
69+ d_end: int, # 特征维度结束
70+ cache_mode: int, # 1=分页 2=每 batch 一个 block
71+) -> torch.Tensor:
72+ """
73+ 从 page cache 中读取 [start_seq, end_seq) × [d_start, d_end) 的数据。
74+ 返回 (end_seq - start_seq, d_end - d_start)。
75+ """
76+ block_size = state.shape[1]
77+ n = end_seq - start_seq
78+ chunks: list[torch.Tensor] = []
79+ done = 0
80+ while done < n:
81+ sid = start_seq + done
82+ if cache_mode == 1:
83+ bid = int(block_table[b_idx, sid // block_size].item())
84+ else:
85+ bid = int(block_table[b_idx].item())
86+ off = sid % block_size
87+ cnt = block_size - off
88+ if cnt > n - done:
89+ cnt = n - done
90+ chunks.append(state[bid:bid + 1, off:off + cnt, d_start:d_end].reshape(cnt, d_end - d_start))
91+ done += cnt
92+ return torch.cat(chunks, dim=0)
93+ 
94+ 
95+def _write_cache(
96+ state: torch.Tensor,
97+ block_table: torch.Tensor,
98+ data: torch.Tensor, # (seq_len, coff * head_dim)
99+ b_idx: int,
100+ start_seq: int,
101+ end_seq: int,
102+ cache_mode: int,
103+ update_mask: Optional[torch.Tensor] = None, # 与 state 同 shape 的 bool 张量，记录写入位置
104+) -> torch.Tensor:
105+ """
106+ 向 page cache 写入数据。为避免破坏 autograd 计算图，返回新的 state 张量。
107+ 传入 update_mask 时同步标记实际写入的位置（block_id 为 0 的无效块不标记）。
108+ """
109+ state = state.clone()
110+ block_size = state.shape[1]
111+ n = end_seq - start_seq
112+ done = 0
113+ while done < n:
114+ sid = start_seq + done
115+ if cache_mode == 1:
116+ bid = int(block_table[b_idx, sid // block_size].item())
117+ if bid == 0:
118+ # block_id 为 0 表示无效 block，跳过
119+ done += block_size - (sid % block_size)
120+ continue
121+ else:
122+ bid = int(block_table[b_idx].item())
123+ off = sid % block_size
124+ cnt = block_size - off
125+ if cnt > n - done:
126+ cnt = n - done
127+ state[bid:bid + 1, off:off + cnt, :] = data[done:done + cnt]
128+ if update_mask is not None:
129+ update_mask[bid:bid + 1, off:off + cnt, :] = True
130+ done += cnt
131+ return state
132+ 
133+ 
134+def _build_ape_for_flat_sequence(
135+ total_len: int,
136+ B: int,
137+ start_pos: List[int],
138+ ape: torch.Tensor,
139+ cmp_ratio: int,
140+ bs_combine: bool,
141+ cu_seqlens: Optional[torch.Tensor],
142+) -> torch.Tensor:
143+ """
144+ 为展平序列的每个位置预计算 ape。
145+ 返回 (total_len, coff*head_dim)，
146+ 其中第 p 行 = ape[全局序列位置 % cmp_ratio, :]。
147+ """
148+ ape_out = torch.zeros(total_len, ape.shape[1], dtype=torch.float32, device=ape.device)
149+ for b in range(B):
150+ sp = start_pos[b]
151+ if bs_combine:
152+ fs = int(cu_seqlens[b].item())
153+ fe = int(cu_seqlens[b + 1].item())
154+ else:
155+ fs = b * (total_len // B) if B > 0 else 0
156+ fe = (b + 1) * (total_len // B) if B > 0 else total_len
157+ L = fe - fs
158+ if L <= 0:
159+ continue
160+ idx = (torch.arange(sp, sp + L) % cmp_ratio).long()
161+ ape_out[fs:fe, :] = ape[idx, :]
162+ return ape_out
163+ 
164+ 
165+# ===========================================================================
166+# 前向传播（纯 PyTorch，autograd 自动处理反向）
167+# ===========================================================================
168+ 
169+def compressor_forward(
170+ x: torch.Tensor, # [B, S, H] 或 [T, H]（bs_combine 模式）
171+ wkv: torch.Tensor, # (coff * head_dim, H)
172+ wgate: torch.Tensor, # (coff * head_dim, H)
173+ kv_state: torch.Tensor, # (block_num, block_size, coff * head_dim)
174+ score_state: torch.Tensor, # (block_num, block_size, coff * head_dim)
175+ ape: torch.Tensor, # (cmp_ratio, coff * head_dim) 位置编码
176+ block_table: torch.Tensor, # (B, max_blocks) int32 page 映射表
177+ start_pos: List[int], # 每个 batch 的起始序列位置
178+ cmp_ratio: int = 4, # 压缩率：每 cmp_ratio 个 token 压缩为 1 个
179+ coff: int = 1, # 1=无重叠 2=相邻块重叠
180+ cache_mode: int = 1, # 1=分页 cache 2=每 batch 一个 block
181+ cu_seqlens: Optional[torch.Tensor] = None, # bs_combine 模式下的累积序列长度
182+ seqused: Optional[List[int]] = None, # 每个 batch 实际使用的序列长度
183+ return_intermediates: bool = False, # 是否返回 d_kv, d_score, d_sc2
184+ return_update_mask: bool = False, # 是否返回 kv/score state 的写入位置 mask（与 kv_state_out/score_state_out 同 shape）
185+) -> Tuple:
186+ """
187+ 执行 compressor 前向计算。
188+ 
189+ 返回:
190+ cmp_kv: 压缩后的 KV，(压缩块数, head_dim) 或 (B, 最大块数, head_dim)
191+ cmp_kv_mask: 同形状 bool 张量，标记有效位置
192+ kv_state: 更新后的 kv_state
193+ score_state: 更新后的 score_state
194+ """
195+ orig_dtype = x.dtype
196+ head_dim = wkv.shape[0] // coff
197+ B = len(start_pos)
198+ bs_combine = (cu_seqlens is not None)
199+ 
200+ # ---- 第一步：matmul ----
201+ x_f32 = x.float()
202+ wkv_f32 = wkv.float()
203+ wgate_f32 = wgate.float()
204+ ape_f32 = ape.float()
205+ 
206+ if bs_combine:
207+ # x 形状为 (T, H)，直接做矩阵乘法
208+ flat_kv = x_f32 @ wkv_f32.T # (T, coff*head_dim)
209+ flat_score = x_f32 @ wgate_f32.T
210+ T = x.shape[0]
211+ else:
212+ # x 形状为 (B, S, H)，先展平再乘
213+ B_in, S_in, H_in = x.shape
214+ flat_kv = x_f32.reshape(B_in * S_in, H_in) @ wkv_f32.T
215+ flat_score = x_f32.reshape(B_in * S_in, H_in) @ wgate_f32.T
216+ T = B_in * S_in
217+ 
218+ # 模拟 kernel 反向中间精度：dkv/dsb 梯度经 FP16/BF16 存储（golden 同款量化）
219+ flat_kv = _QuantGrad.apply(flat_kv, orig_dtype)
220+ flat_score = _QuantGrad.apply(flat_score, orig_dtype)
221+ # ---- 第二步：全局添加 ape（与 golden 中 in-place 加法等价）----
222+ ape_add = _build_ape_for_flat_sequence(T, B, start_pos, ape_f32, cmp_ratio, bs_combine, cu_seqlens)
223+ flat_score = flat_score + ape_add # (T, coff*head_dim)，ape 已应用
224+ 
225+ # ---- 第三步：分配输出缓冲区 ----
226+ # 与 golden 的对齐策略：TH 布局时预分配 (min(T, T//cmp_ratio+B), head_dim)
227+ # BSH 布局时预分配 (B, ceil(S/cmp_ratio), head_dim)
228+ if not bs_combine:
229+ S_val = x.shape[1]
230+ max_blocks = (S_val + cmp_ratio - 1) // cmp_ratio
231+ cmp_kv = torch.zeros(B, max_blocks, head_dim, dtype=torch.float32, device=x.device)
232+ cmp_mask = torch.zeros(B, max_blocks, head_dim, dtype=torch.bool, device=x.device)
233+ else:
234+ T_val = x.shape[0]
235+ out_rows = min(T_val, T_val // cmp_ratio + B)
236+ cmp_kv = torch.zeros(out_rows, head_dim, dtype=torch.float32, device=x.device)
237+ cmp_mask = torch.zeros(out_rows, head_dim, dtype=torch.bool, device=x.device)
238+ out_idx_flat = 0 # 当前写入行索引
239+ 
240+ kv_state_out = kv_state.clone()
241+ score_state_out = score_state.clone()
242+ kv_update_mask = torch.zeros_like(kv_state_out, dtype=torch.bool)
243+ score_update_mask = torch.zeros_like(score_state_out, dtype=torch.bool)
244+ 
245+ # 中间变量：按 output layout 预分配（与 cmp_kv 同 shape，最后维为 N=coff*cmpRatio）
246+ N_inter = coff * cmp_ratio
247+ if return_intermediates:
248+ if not bs_combine:
249+ kvIntermediate = torch.zeros(B, max_blocks, N_inter, head_dim, dtype=torch.float32, device=x.device)
250+ softmaxScoreIntermediate = torch.zeros(B, max_blocks, N_inter, head_dim, dtype=torch.float32, device=x.device)
251+ else:
252+ kvIntermediate = torch.zeros(out_rows, N_inter, head_dim, dtype=torch.float32, device=x.device)
253+ softmaxScoreIntermediate = torch.zeros(out_rows, N_inter, head_dim, dtype=torch.float32, device=x.device)
254+ 
255+ if bs_combine:
256+ out_idx_flat = 0
257+ 
258+ # ---- 第四步：逐 batch 逐块处理 ----
259+ for b_idx in range(B):
260+ b_start = start_pos[b_idx]
261+ 
262+ # 确定当前 batch 的序列长度
263+ if seqused is not None:
264+ b_seq_used = seqused[b_idx]
265+ elif bs_combine:
266+ b_seq_used = int(cu_seqlens[b_idx + 1].item() - cu_seqlens[b_idx].item())
267+ else:
268+ b_seq_used = x.shape[1]
269+ 
270+ # compress_seq_id：小于此位置的需要压缩
271+ compress_limit = (b_start + b_seq_used) // cmp_ratio * cmp_ratio
272+ batch_out_sc_id = 0
273+ batch_seq_idx = 0
274+ 
275+ while batch_seq_idx < b_seq_used:
276+ # 当前块的全局起始/结束序列位置
277+ s_start = b_start + batch_seq_idx
278+ s_end = s_start // cmp_ratio * cmp_ratio + cmp_ratio
279+ if s_end > b_start + b_seq_used:
280+ s_end = b_start + b_seq_used
281+ 
282+ # 展平数组中的偏移
283+ base = int(cu_seqlens[b_idx].item()) if bs_combine else b_idx * x.shape[1]
284+ off_s = base + (s_start - b_start)
285+ off_e = base + (s_end - b_start)
286+ 
287+ # 是否保存到 cache / 是否压缩
288+ save_flag = (cache_mode == 1) or (
289+ s_start >= (compress_limit - (coff - 1) * cmp_ratio)
290+ )
291+ compress_flag = (s_start < compress_limit)
292+ 
293+ # ---- 写入 page cache（kv 无 ape，score 已带 ape）----
294+ if save_flag:
295+ kv_state_out = _write_cache(kv_state_out, block_table,
296+ flat_kv[off_s:off_e, :], b_idx,
297+ s_start, s_end, cache_mode,
298+ kv_update_mask if return_update_mask else None)
299+ score_state_out = _write_cache(score_state_out, block_table,
300+ flat_score[off_s:off_e, :], b_idx,
301+ s_start, s_end, cache_mode,
302+ score_update_mask if return_update_mask else None)
303+ 
304+ # ---- 执行压缩 ----
305+ if compress_flag:
306+ # 初始化 sc 缓冲区
307+ kvLocal = torch.zeros(coff, cmp_ratio, head_dim, device=x.device)
308+ sc_score = torch.full((coff, cmp_ratio, head_dim), float('-inf'), device=x.device)
309+ 
310+ # --- 填充当前数据（coff_id = coff-1）---
311+ cur_cid = coff - 1
312+ ds_cur = cur_cid * head_dim
313+ de_cur = (cur_cid + 1) * head_dim
314+ cfs = 0 # 从 state cache 读取的数量
315+ if b_start == s_start:
316+ cfs = b_start % cmp_ratio
317+ if cfs > 0:
318+ kvLocal[cur_cid, 0:cfs, :] = _read_cache(
319+ kv_state_out, block_table, b_idx,
320+ b_start - cfs, b_start, ds_cur, de_cur, cache_mode)
321+ sc_score[cur_cid, 0:cfs, :] = _read_cache(
322+ score_state_out, block_table, b_idx,
323+ b_start - cfs, b_start, ds_cur, de_cur, cache_mode)
324+ 
325+ # 从展平数组获取当前数据（kv 无 ape，score 已预加 ape）
326+ kvLocal[cur_cid, cfs:cmp_ratio, :] = flat_kv[off_s:off_e, ds_cur:de_cur]
327+ sc_score[cur_cid, cfs:cmp_ratio, :] = flat_score[off_s:off_e, ds_cur:de_cur]
328+ 
329+ # --- 填充前一块数据（coff_id = 0），仅 coff == 2 时有效 ---
330+ if coff == 2:
331+ prev_cid = 0
332+ ds_pr = prev_cid * head_dim
333+ de_pr = (prev_cid + 1) * head_dim
334+ cfs_p = 0
335+ if b_start == s_start:
336+ # 第一个块：前一块数据全部来自 cache
337+ cfs_p = cmp_ratio
338+ if b_start >= cmp_ratio:
339+ cs2 = b_start - b_start % cmp_ratio - cmp_ratio
340+ kvLocal[prev_cid, 0:cfs_p, :] = _read_cache(
341+ kv_state_out, block_table, b_idx,
342+ cs2, cs2 + cfs_p, ds_pr, de_pr, cache_mode)
343+ sc_score[prev_cid, 0:cfs_p, :] = _read_cache(
344+ score_state_out, block_table, b_idx,
345+ cs2, cs2 + cfs_p, ds_pr, de_pr, cache_mode)
346+ elif s_start - cmp_ratio < b_start:
347+ # 第二个块：部分来自 cache，部分来自当前数据的前半段
348+ cfs_p = b_start % cmp_ratio
349+ if cfs_p > 0:
350+ cs2 = b_start - cfs_p
351+ kvLocal[prev_cid, 0:cfs_p, :] = _read_cache(
352+ kv_state_out, block_table, b_idx,
353+ cs2, b_start, ds_pr, de_pr, cache_mode)
354+ sc_score[prev_cid, 0:cfs_p, :] = _read_cache(
355+ score_state_out, block_table, b_idx,
356+ cs2, b_start, ds_pr, de_pr, cache_mode)
357+ if cfs_p < cmp_ratio:
358+ po_s = off_s - (cmp_ratio - cfs_p)
359+ po_e = off_s
360+ kvLocal[prev_cid, cfs_p:cmp_ratio, :] = flat_kv[po_s:po_e, ds_pr:de_pr]
361+ sc_score[prev_cid, cfs_p:cmp_ratio, :] = flat_score[po_s:po_e, ds_pr:de_pr]
362+ 
363+ # ---- 核心压缩：softmax + 加权求和 ----
364+ kv_2d = kvLocal.reshape(coff * cmp_ratio, head_dim)
365+ sc_score_2d = sc_score.reshape(coff * cmp_ratio, head_dim)
366+ # 将 -inf 替换为有限值以保证 softmax 正确
367+ sc_score_2d = torch.where(torch.isinf(sc_score_2d),
368+ torch.tensor(-1e9, device=x.device), sc_score_2d)
369+ # 按列 softmax → 权重
370+ sm = torch.softmax(sc_score_2d, dim=0)
371+ # 逐元素乘 + 沿 token 维求和
372+ compressed = (kv_2d * sm).sum(dim=0, keepdim=True) # (1, head_dim)
373+ 
374+ if return_intermediates:
375+ if not bs_combine:
376+ kvIntermediate[b_idx, batch_out_sc_id] = kv_2d.detach()
377+ softmaxScoreIntermediate[b_idx, batch_out_sc_id] = sm.detach()
378+ else:
379+ kvIntermediate[out_idx_flat] = kv_2d.detach()
380+ softmaxScoreIntermediate[out_idx_flat] = sm.detach()
381+ 
382+ if not bs_combine:
383+ cmp_kv[b_idx, batch_out_sc_id, :] = compressed
384+ cmp_mask[b_idx, batch_out_sc_id, :] = True
385+ else:
386+ cmp_kv[out_idx_flat, :] = compressed
387+ cmp_mask[out_idx_flat, :] = True
388+ out_idx_flat += 1
389+ 
390+ batch_out_sc_id += 1
391+ 
392+ batch_seq_idx = s_end - b_start
393+ 
394+ # ---- 组装最终输出 ----
395+ # TH 和 BSH 都使用预分配缓冲区，无需额外拼接
396+ 
397+ cmp_kv = cmp_kv.to(orig_dtype)
398+ if return_intermediates and return_update_mask:
399+ return (cmp_kv, cmp_mask, kv_state_out, score_state_out,
400+ softmaxScoreIntermediate, kvIntermediate,
401+ kv_update_mask, score_update_mask)
402+ if return_update_mask:
403+ return (cmp_kv, cmp_mask, kv_state_out, score_state_out,
404+ kv_update_mask, score_update_mask)
405+ if return_intermediates:
406+ return cmp_kv, cmp_mask, kv_state_out, score_state_out, softmaxScoreIntermediate, kvIntermediate
407+ return cmp_kv, cmp_mask, kv_state_out, score_state_out
408+ 
409+ 
410+# ===========================================================================
411+# 便捷包装：可直接在模型中调用的 Compressor 模块
412+# ===========================================================================
413+ 
414+class Compressor(torch.nn.Module):
415+ """
416+ Compressor 算子模块，支持前向 + 反向。
417+ 
418+ 用法:
419+ compressor = Compressor(cmp_ratio=4, coff=1, cache_mode=1)
420+ cmp_kv, cmp_mask, kv_state, score_state = compressor(
421+ x, wkv, wgate, kv_state, score_state, ape,
422+ block_table, start_pos, cu_seqlens, seqused)
423+ """
424+ 
425+ def __init__(self, cmp_ratio: int = 4, coff: int = 1, cache_mode: int = 1):
426+ super().__init__()
427+ self.cmp_ratio = cmp_ratio
428+ self.coff = coff
429+ self.cache_mode = cache_mode
430+ 
431+ def forward(
432+ self,
433+ x: torch.Tensor,
434+ wkv: torch.Tensor,
435+ wgate: torch.Tensor,
436+ kv_state: torch.Tensor,
437+ score_state: torch.Tensor,
438+ ape: torch.Tensor,
439+ block_table: torch.Tensor,
440+ start_pos: List[int],
441+ cu_seqlens: Optional[torch.Tensor] = None,
442+ seqused: Optional[List[int]] = None,
443+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
444+ return compressor_forward(
445+ x, wkv, wgate, kv_state, score_state, ape,
446+ block_table, start_pos,
447+ cmp_ratio=self.cmp_ratio, coff=self.coff, cache_mode=self.cache_mode,
448+ cu_seqlens=cu_seqlens, seqused=seqused,
449+ )
Aattention/compressor_grad/tests/pytest/compressor_grad_check.py+413-0
@@ -0,0 +1,413 @@
1+#!/usr/bin/python
2+# -*- coding: utf-8 -*-
3+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+# See LICENSE in the root of the software repository for the full text of the License.
10+# ======================================================================================================================
11+ 
12+import numpy as np
13+import datetime
14+import os
15+import sys
16+import logging
17+ 
18+ 
19+np.set_printoptions(suppress=True)
20+ 
21+logging.basicConfig(level=logging.INFO, format='%(message)s', force=True)
22+logger = logging.getLogger(__name__)
23+ 
24+ 
25+# ================================================================
26+# 精度阈值配置
27+# ================================================================
28+# rtol/atol: np.isclose 的相对/绝对容差
29+# pct_thd: 通过率阈值，0.005 表示 99.5% 的点位需满足 rtol/atol
30+# diff_thd: 相对误差归一化分母系数（b = max(|real|,|expect|, 1/(2^14)/diff_thd)）
31+# max_diff_hd: 单点最大允许相对误差阈值，超过即判 Failed
32+PRECISION = {
33+ "bfloat16": {"rtol": 0.0078125, "atol": 0.0001, "pct_thd": 0.005,
34+ "diff_thd": 0.005, "max_diff_hd": 10.0},
35+ "float16": {"rtol": 0.005, "atol": 0.000025, "pct_thd": 0.005,
36+ "diff_thd": 0.005, "max_diff_hd": 10.0},
37+ "float32": {"rtol": 0.005, "atol": 0.000025, "pct_thd": 0.005,
38+ "diff_thd": 0.005, "max_diff_hd": 10.0},
39+}
40+ 
41+ 
42+# ================================================================
43+# 输出定义
44+# ================================================================
45+# 算子所有输出列表（顺序固定，便于汇总统计）：
46+# 反向 4 个 + 正向 3 个 + state_cache 4 个；通路 2/3 的 result 无正向/state 字段，
47+# 打印/汇总时自动跳过
48+ALL_OUTPUTS = ["d_wkv", "d_wgate", "d_ape", "d_x",
49+ "cmp_kv", "softmax_score", "kv",
50+ "kv_state_update", "score_state_update",
51+ "kv_state_origin", "score_state_origin"]
52+ 
53+# 输出名 → result dict 中的 status/pct key 映射
54+_OUTPUT_KEY_MAP = {
55+ "d_wkv": ("dwkvPct", "dwkvStatus"),
56+ "d_wgate": ("dwgatePct", "dwgateStatus"),
57+ "d_ape": ("apePct", "apeStatus"),
58+ "d_x": ("dxPct", "dxStatus"),
59+ "cmp_kv": ("cmpkvPct", "cmpkvStatus"),
60+ "softmax_score": ("smPct", "smStatus"),
61+ "kv": ("kvPct", "kvStatus"),
62+ "kv_state_update": ("kvUpdPct", "kvUpdStatus"),
63+ "score_state_update": ("scoreUpdPct", "scoreUpdStatus"),
64+ "kv_state_origin": ("kvOrgPct", "kvOrgStatus"),
65+ "score_state_origin": ("scoreOrgPct", "scoreOrgStatus"),
66+}
67+ 
68+ 
69+def get_output_keys(out_name):
70+ """返回输出名对应的 (pct_key, status_key)。"""
71+ return _OUTPUT_KEY_MAP[out_name]
72+ 
73+ 
74+def get_pct_thd(data_type):
75+ """根据 dtype 字符串返回通过率阈值 pct_thd。"""
76+ return PRECISION[data_type]["pct_thd"]
77+ 
78+ 
79+# ================================================================
80+# 日志输出
81+# ================================================================
82+def print_log(data=None, level='INFO'):
83+ stamp = datetime.datetime.now().strftime("%Y/%m/%d %H:%M:%S")
84+ caller = sys._getframe().f_back
85+ caller_file = os.path.basename(caller.f_code.co_filename)
86+ caller_line = str(caller.f_lineno).zfill(4)
87+ print(f"[{stamp}] [{level}]-{caller_file}:{caller_line} - {data}")
88+ 
89+ 
90+# ================================================================
91+# 单点位精度对比展示
92+# ================================================================
93+# 精度对比表格的共用分隔线与表头（display/check 系列函数共享）
94+_ROW_DIVIDER = '-' * 87
95+_ERROR_DIVIDER = 'Error Line' + '-' * 77
96+_MAX_RE_DIVIDER = 'Max-RE line:' + '-' * 75
97+_LOOP_HEADER = 'Loop \t ExpectOut \t RealOut \t FpDiff \t RateDiff'
98+ 
99+ 
100+def cal_relative_diff_np_isclose(real_data, expect_data):
101+ """单点相对误差：|real - expect| / (|expect| + 1e-9)。"""
102+ diff = abs(float(real_data) - float(expect_data))
103+ return diff / (abs(float(expect_data)) + 1e-9)
104+ 
105+ 
106+def _format_loop_row(seq, expect_val, real_val):
107+ """格式化单行点位：期望值为 inf/nan 时用字符串列，否则用 7 位小数。"""
108+ diff_rate = cal_relative_diff_np_isclose(real_val, expect_val)
109+ if "inf" in str(expect_val) or "nan" in str(expect_val):
110+ diff_abs = "inf" if "inf" in str(expect_val) else "nan"
111+ return f"{seq:08d} \t {expect_val:<7} \t {real_val:<7} \t {diff_abs:<7} \t {diff_rate:<7}"
112+ diff_abs = abs(np.float64(expect_val) - np.float64(real_val))
113+ return f"{seq:08d} \t {expect_val:.7f} \t {real_val:.7f} \t {diff_abs:.7f} \t {diff_rate:.7f}"
114+ 
115+ 
116+def _format_error_row(seq, expect_val, real_val, diff_rate):
117+ """格式化错误点位明细行（7 位小数定点）。"""
118+ diff_abs = abs(np.float64(expect_val) - np.float64(real_val))
119+ return f"{seq:08d} \t {expect_val:.7f} \t {real_val:.7f} \t {diff_abs:.7f} \t {diff_rate:.7f}"
120+ 
121+ 
122+def _display_row_range(real_data, expect_data, start, row_range):
123+ """打印 [start + lo, start + hi) 区间的点位行（绝对序号 = start + idx + 1）。"""
124+ for idx in row_range:
125+ seq = start + idx + 1
126+ print_log(_format_loop_row(seq, expect_data[idx + start], real_data[idx + start]))
127+ 
128+ 
129+def display_output_np_isclose(real_data, expect_data, start, end):
130+ """输出每个点位的对比结果（≤20 点位全显示，>20 显示首尾各 10）。"""
131+ print_log(_ROW_DIVIDER)
132+ print_log(_LOOP_HEADER)
133+ print_log(_ROW_DIVIDER)
134+ split_count = int(end - start)
135+ if split_count <= 20:
136+ _display_row_range(real_data, expect_data, start, range(split_count + 1))
137+ else:
138+ _display_row_range(real_data, expect_data, start, range(10))
139+ print_log('... \t ... \t ... \t ... \t ...')
140+ _display_row_range(real_data, expect_data, start,
141+ range(split_count - 10 + 1, split_count + 1))
142+ 
143+ 
144+def display_error_output(real_data, expect_data, err_idx, relative_diff):
145+ """输出错误点位明细 + 最大相对误差点位。"""
146+ print_log(_ERROR_DIVIDER)
147+ print_log(_LOOP_HEADER)
148+ print_log(_ROW_DIVIDER)
149+ err_num = len(err_idx)
150+ for count, i in enumerate(err_idx, start=1):
151+ if count < 10 or (90 < count < 100):
152+ print_log(_format_error_row(i, expect_data[i], real_data[i], relative_diff[count - 1]))
153+ elif count == 10 or (count == 100 and err_num > 100):
154+ print_log(f"{'...':>8} \t {'...':>7} \t {'...':>7} \t {'...':>7} \t {'...':>7}")
155+ elif count > 100:
156+ break
157+ 
158+ print_log(_MAX_RE_DIVIDER)
159+ max_error = max(relative_diff)
160+ m_idx_list = err_idx[np.where(relative_diff == max_error)]
161+ for m_idx in m_idx_list[:4]:
162+ print_log(_format_error_row(m_idx, expect_data[m_idx], real_data[m_idx], max_error))
163+ print_log(_ROW_DIVIDER)
164+ 
165+ 
166+def _print_verdict(rtol, atol, pct_thd, fulfill_percent, result_str):
167+ """打印 Rtol/Atol/PctThd/PctRlt/Result 汇总表格（check_result 两分支共用）。"""
168+ print_log(_ROW_DIVIDER)
169+ print_log('Rtol \t Atol \t PctThd \t PctRlt \t Result')
170+ print_log(_ROW_DIVIDER)
171+ print_log(f"{rtol:.4f} \t {atol:.6f} \t {pct_thd:.2f}% \t {fulfill_percent:.6f}% \t {result_str}")
172+ 
173+ 
174+# ================================================================
175+# 单输出精度对比主函数
176+# ================================================================
177+def check_result(expect, result, data_type, pct_thd=0.005):
178+ """单输出精度对比，复用自正向 compressor_golden.check_result。
179+ 
180+ Args:
181+ expect: CPU golden 结果 (torch.Tensor)
182+ result: NPU 反向算子输出 (torch.Tensor)
183+ data_type: dtype 字符串 ('bfloat16' / 'float16' / 'float32')
184+ pct_thd: 通过率阈值（默认 0.005 表示 99.5%）
185+ 
186+ Returns:
187+ fulfill_percent: 满足精度要求的百分比
188+ result_str: "Pass" / "Failed"
189+ """
190+ real_data = result.cpu().numpy().flatten()
191+ data_compe = expect.cpu().numpy().flatten()
192+ if real_data.size == 0 and data_compe.size == 0:
193+ print_log('The npu_output is [],and it is same as bm_output, the result of data_compare is "Pass"')
194+ return 100.0, "Pass"
195+ max_error = 0
196+ result_str = "Failed"
197+ start, end = 0, max(real_data.size - 1, 0)
198+ 
199+ if real_data.size != data_compe.size:
200+ print_log(f"Error,the size of npu output[{real_data.size}] and benchmark[{data_compe.size}] is not equal.")
201+ return 0.0, result_str
202+ overflows_count = data_compe[np.isinf(data_compe)].size + data_compe[np.isnan(data_compe)].size
203+ 
204+ if overflows_count > 0:
205+ print_log(f"Overflow,size:{overflows_count},benchmark_output:"
206+ f"{data_compe[np.isinf(data_compe)][0:10]}, {data_compe[np.isnan(data_compe)][0:10]}")
207+ 
208+ # 检测 NPU 输出 (real_data) 中的 NaN/Inf；同时计算“两边同坏”（golden 与 NPU
209+ # 同位置都是 NaN 或都是 Inf）——同坏视为一致（输入语义一致时输出应一致），
210+ # 不判错；单边坏（NPU 坏但 golden 对应位置不坏）才是真实 bug，强制 Failed。
211+ real_data_f32 = real_data.astype(np.float32) if str(real_data.dtype) == 'bfloat16' else real_data
212+ data_compe_f32 = data_compe.astype(np.float32) if str(data_compe.dtype) == 'bfloat16' else data_compe
213+ real_nan_mask = np.isnan(real_data_f32)
214+ real_inf_mask = np.isinf(real_data_f32)
215+ real_overflow_mask = real_nan_mask | real_inf_mask
216+ real_overflow_count = int(real_overflow_mask.sum())
217+ both_bad = (np.isnan(real_data_f32) & np.isnan(data_compe_f32)) | \
218+ (np.isinf(real_data_f32) & np.isinf(data_compe_f32))
219+ unmatched_bad = real_overflow_mask & ~both_bad
220+ unmatched_bad_count = int(unmatched_bad.sum())
221+ if real_overflow_count > 0:
222+ print_log(f"NPU output has NaN/Inf, count:{real_overflow_count}, nan:{int(real_nan_mask.sum())}, "
223+ f"inf:{int(real_inf_mask.sum())}, sample_values:{real_data_f32[real_overflow_mask][0:10]}")
224+ 
225+ # 仅支持 bfloat16 / float16 / float32，阈值统一从 PRECISION 配置获取
226+ cfg = PRECISION[data_type]
227+ diff_thd = cfg["diff_thd"]
228+ max_diff_hd = cfg["max_diff_hd"]
229+ rtol = cfg["rtol"]
230+ atol = cfg["atol"]
231+ max_error_idx = 10000000
232+ 
233+ split_count = int(end - start) + 1
234+ print_log(f"split_count:{float(split_count)}; max_diff_hd:{max_diff_hd};")
235+ 
236+ # bfloat16 需转 float32 再 isclose；float16 / float32 直接 isclose
237+ # 注意：使用 equal_nan=False，使 NPU 输出中的 NaN 与 CPU 期望不匹配时判为错误；
238+ # 但两边同坏（both_bad：同位置都是 NaN 或都是 Inf）视为匹配
239+ if str(real_data.dtype) == 'bfloat16':
240+ diff_result = np.isclose(real_data.astype(np.float32), data_compe.astype(np.float32),
241+ rtol=rtol, atol=atol, equal_nan=False)
242+ else:
243+ diff_result = np.isclose(real_data, data_compe, rtol=rtol, atol=atol, equal_nan=False)
244+ diff_result = diff_result | both_bad
245+ err_idx = np.where(diff_result != np.array((True,)))[0]
246+ 
247+ if data_compe.dtype == np.bool_:
248+ data_compe = data_compe.astype(np.int8)
249+ real_data = real_data.astype(np.int8)
250+ diff_abs = np.abs(data_compe - real_data)
251+ b1 = np.maximum(np.abs(real_data), np.abs(data_compe))
252+ b2 = (1.0 / (1 << 14)) / diff_thd
253+ b = np.maximum(b1, b2) + 1e-9
254+ err_diff = diff_abs / (b + 1e-9)
255+ err_diff = err_diff[err_idx]
256+ 
257+ fulfill_percent = float(split_count - err_idx.size) / float(split_count) * 100.0
258+ 
259+ # NPU 输出含 NaN/Inf 且 golden 对应位置不同样坏（unmatched）时判 Failed
260+ #（NaN 会传染，结果不可信）；两边同坏（both_bad）视为一致，不强制 Failed
261+ if unmatched_bad_count > 0:
262+ result_str = "Failed"
263+ max_error = float('inf')
264+ display_output_np_isclose(real_data, data_compe, start, end)
265+ pct_thd = (1 - pct_thd) * 100.0
266+ _print_verdict(rtol, atol, pct_thd, fulfill_percent, result_str)
267+ print_log(f"NPU output has NaN/Inf unmatched with golden (count={unmatched_bad_count}, "
268+ f"total nan/inf={real_overflow_count}). Force Failed.")
269+ if len(err_diff) > 0:
270+ display_error_output(real_data, data_compe, err_idx, err_diff[0:max_error_idx])
271+ return fulfill_percent, result_str
272+ 
273+ display_output_np_isclose(real_data, data_compe, start, end)
274+ pct_thd = (1 - pct_thd) * 100.0
275+ result_str = "Pass" if (fulfill_percent >= pct_thd) else "Failed"
276+ if len(err_diff) > 0:
277+ # 过滤 NaN/Inf 后再取 max，避免 NaN 导致 max_error 比较失效
278+ finite_err_diff = err_diff[np.isfinite(err_diff)]
279+ if len(finite_err_diff) > 0:
280+ max_error = float(np.max(finite_err_diff[0:max_error_idx]))
281+ if max_error >= max_diff_hd:
282+ result_str = "Failed"
283+ _print_verdict(rtol, atol, pct_thd, fulfill_percent, result_str)
284+ if len(err_diff) > 0:
285+ print_log(f"Max-RelativeError is: {max_error}. Threshold is: {max_diff_hd}.")
286+ if result_str == "Failed":
287+ display_error_output(real_data, data_compe, err_idx, err_diff[0:max_error_idx])
288+ return fulfill_percent, result_str
289+ 
290+ 
291+def check_one_output(name, expect, result, data_type, enabled, total_valid, pct_thd=0.005):
292+ """对比单个反向输出，未开启对比的返回 SKIP 标记。
293+ 
294+ Args:
295+ name: 输出名 (如 'd_ape', 'd_x', 'd_wkv', 'd_wgate')
296+ expect: CPU golden 结果 (torch.Tensor)
297+ result: NPU 反向算子输出 (torch.Tensor)
298+ data_type: dtype 字符串 ('bfloat16' / 'float16' / 'float32')
299+ enabled: 是否进行精度对比
300+ total_valid: 该用例的有效压缩块数（0 时无有效数据，自动 SKIP）
301+ pct_thd: 通过率阈值
302+ 
303+ Returns:
304+ dict: {diff, pct, status}
305+ """
306+ if not enabled:
307+ return {"diff": float("nan"), "pct": 100.0, "status": "SKIP"}
308+ 
309+ print_log("=" * 80)
310+ print_log(f"check {name}:")
311+ print_log("--------------------------------------------------------------")
312+ try:
313+ if expect is None or result is None:
314+ print_log(f"{name}: expect or result is None, skip")
315+ return {"diff": float("nan"), "pct": 100.0, "status": "SKIP"}
316+ if expect.shape != result.shape:
317+ print_log(f"{name}: shape mismatch expect={expect.shape} result={result.shape}, flatten compare")
318+ fulfill_percent, result_str = check_result(expect, result, data_type, pct_thd)
319+ diff = expect.cpu().float().flatten() - result.cpu().float().flatten()
320+ max_abs_diff = diff.abs().max().item() if diff.numel() > 0 else 0.0
321+ return {"diff": max_abs_diff, "pct": fulfill_percent, "status": "PASS" if result_str == "Pass" else "FAIL"}
322+ except Exception as e:
323+ import traceback
324+ traceback.print_exc()
325+ return {"diff": float("nan"), "pct": 0.0, "status": "ERROR"}
326+ 
327+ 
328+# ================================================================
329+# 结果格式化与汇总统计
330+# ================================================================
331+def format_output_status(result, out_name):
332+ """格式化单个输出的状态字符串（正反向统一格式，用于用例行展示）。
333+ 
334+ 输出格式（正反向完全一致）：`<名>=<pct>%[<STATUS>]`；
335+ result 中不存在该输出的 status 字段（如通路 2/3 无正向字段）时返回 None，
336+ 由调用方跳过，不参与展示。
337+ """
338+ pct_key, stat_key = get_output_keys(out_name)
339+ if stat_key not in result:
340+ return None
341+ st = result[stat_key]
342+ pct = result.get(pct_key, 100.0)
343+ if st == "SKIP":
344+ return f"{out_name}=SKIP"
345+ elif st == "ERROR":
346+ return f"{out_name}=ERROR"
347+ else:
348+ return f"{out_name}={pct:.2f}%[{st}]"
349+ 
350+ 
351+def format_case_line(result):
352+ """格式化单个用例的完整结果行（正反向统一；按 result 实际含有的输出展示）。"""
353+ parts = []
354+ for out_name in ALL_OUTPUTS:
355+ s = format_output_status(result, out_name)
356+ if s is not None:
357+ parts.append(s)
358+ detail = " ".join(parts)
359+ return f" {result.get('status', 'N/A'):5s} {detail} (validBlocks={result.get('totalValid', 0)})"
360+ 
361+ 
362+def build_error_result(name):
363+ """构建异常用例的 result dict，所有输出标记为 ERROR。"""
364+ err_result = dict(name=name, status="ERROR")
365+ for out_name in ALL_OUTPUTS:
366+ _, stat_key = get_output_keys(out_name)
367+ err_result[stat_key] = "ERROR"
368+ return err_result
369+ 
370+ 
371+def print_summary(results):
372+ """输出汇总统计：用例通过个数 + 每个用例每个输出的对比结果（以 case 为粒度）。
373+ 
374+ 输出两部分：
375+ 1. SUMMARY: 总体 PASS/FAIL/SKIP/ERROR 个数统计
376+ 2. Per-Case Per-Output Summary: 以 case 为粒度的表格，
377+ 每行一个用例，每列一个输出，明确看出每个 case 的对比结果
378+ """
379+ print("\n" + "=" * 100)
380+ passed = sum(1 for r in results if r.get("status") == "PASS")
381+ failed = sum(1 for r in results if r.get("status") == "FAIL")
382+ skipped = sum(1 for r in results if r.get("status") == "SKIP")
383+ errors = sum(1 for r in results if r.get("status") == "ERROR")
384+ print(f"SUMMARY: {passed} PASS, {failed} FAIL, {skipped} SKIP, "
385+ f"{errors} ERROR out of {len(results)}")
386+ print("=" * 100)
387+ 
388+ # 以 case 为粒度的每个输出对比结果
389+ # 列由该批 result 实际含有的输出决定（数据驱动，不区分正反向）：
390+ # 通路 1 含正向字段 → 7 列；通路 2/3 仅反向 → 4 列
391+ cols = [o for o in ALL_OUTPUTS
392+ if any(get_output_keys(o)[1] in r for r in results)]
393+ print("\nPer-Case Per-Output Summary:")
394+ print("-" * 100)
395+ header = f"{'Case':<45}" + "".join(f"{o:<18}" for o in cols) + "Overall"
396+ print(header)
397+ print("-" * 100)
398+ for r in results:
399+ case_name = r.get("name", "unknown")
400+ row = f"{case_name:<45}"
401+ for out_name in cols:
402+ pct_key, stat_key = get_output_keys(out_name)
403+ st = r.get(stat_key, "SKIP")
404+ if st == "SKIP":
405+ row += f"{'SKIP':<18}"
406+ elif st == "ERROR":
407+ row += f"{'ERROR':<18}"
408+ else:
409+ pct = r.get(pct_key, 0.0)
410+ row += f"{pct:.2f}%[{st}]".ljust(18)
411+ row += r.get("status", "N/A")
412+ print(row)
413+ print("=" * 100)
Aattention/compressor_grad/tests/pytest/compressor_grad_cpu_golden.py+339-0
@@ -0,0 +1,339 @@
1+#!/usr/bin/python
2+# -*- coding: utf-8 -*-
3+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+# See LICENSE in the root of the software repository for the full text of the License.
10+# ======================================================================================================================
11+ 
12+ 
13+import torch
14+from typing import Optional, List, Tuple
15+ 
16+ 
17+def compressor_grad_golden(
18+ x: torch.Tensor,
19+ wkv: torch.Tensor,
20+ wgate: torch.Tensor,
21+ d_cpm_kv: torch.Tensor,
22+ softmax_score: torch.Tensor,
23+ kv: torch.Tensor,
24+ cu_seqlens: Optional[torch.Tensor] = None,
25+ seqused: Optional[torch.Tensor] = None,
26+ start_pos: Optional[torch.Tensor] = None,
27+ cmp_ratio: int = 4,
28+ coff: int = 1,
29+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
30+ """
31+ Compressor 算子的反向传播。
32+ 
33+ 参数:
34+ x: 输入张量 [B, S, H] 或 [T, H]（BSH 布局或 TH 布局）
35+ wkv: KV 投影权重 (coff * head_dim, H)
36+ wgate: Gate 投影权重 (coff * head_dim, H)
37+ d_cpm_kv: 前向输出 compressed KV 的上游梯度 ∂Loss/∂cmp_kv
38+ BSH: (B, ceil(S/cmp_ratio), head_dim)
39+ TH: (min(T, T//cmp_ratio+B), head_dim)
40+ kv: 前向保存的中间变量（reshape 后的 KV 数据）
41+ BSH: (B, ceil(S/cmp_ratio), coff*cmp_ratio, head_dim)
42+ TH: (min(T, T//cmp_ratio+B), coff*cmp_ratio, head_dim)
43+ softmax_score: 前向保存的 softmax 权重
44+ 形状同 kv
45+ start_pos: 每个 batch 的起始全局序列位置 (B,) int32。
46+ None 时默认全为 0（即所有 batch 从头开始）
47+ cmp_ratio: 压缩率：每 cmp_ratio 个 token 压缩为 1 个
48+ coff: 重叠因子（1=无重叠, 2=相邻块有 50% 重叠）
49+ cu_seqlens: TH 布局下的累积序列长度 (B+1,) int32。
50+ BSH 布局时传 None，函数内部不依赖此参数定位 token
51+ seqused: 每个 batch 实际参与压缩的 token 数 (B,) int32。
52+ None 表示使用全部 sequence length
53+ 
54+ 返回:
55+ d_x: 输入 x 的梯度 ∂Loss/∂x
56+ d_wkv: 权重 wkv 的梯度 ∂Loss/∂wkv
57+ d_wgate: 权重 wgate 的梯度 ∂Loss/∂wgate
58+ d_ape: APE 位置编码的梯度 ∂Loss/∂ape, (cmp_ratio, coff * head_dim)
59+ """
60+ # ══════════════════════════════════════════════════════════════════════
61+ # 阶段 A: 解析输入参数，确定数据布局
62+ # ══════════════════════════════════════════════════════════════════════
63+ 
64+ # head_dim: 每个注意力头的维度 D。
65+ # wkv 的 shape 是 (coff * head_dim, H)，所以 head_dim = wkv.shape[0] // coff
66+ head_dim = wkv.shape[0] // coff
67+ 
68+ # 两种数据布局：
69+ # TH 布局 (is_th_layout=True): x shape (T, H)，B 和 S 合并到 T 维
70+ # 此时 cu_seqlens 提供每个 batch 在 T 维的起止边界
71+ # BSH 布局 (is_th_layout=False): x shape (B, S, H)，batch 独立存储
72+ # 此时不需要 cu_seqlens，直接用 b_idx * S 定位
73+ is_th_layout = (cu_seqlens is not None)
74+ 
75+ # B: batch 数量。start_pos 为 None 时默认全 0
76+ if start_pos is None:
77+ if is_th_layout:
78+ num_batches = cu_seqlens.shape[0] - 1
79+ else:
80+ num_batches = x.shape[0]
81+ start_positions = [0] * num_batches
82+ else:
83+ num_batches = start_pos.shape[0]
84+ start_positions = start_pos.tolist()
85+ 
86+ # 前向计算全部使用 float32（与 golden 一致，保证数值精度）
87+ x_f32 = x.float()
88+ wkv_f32 = wkv.float()
89+ wgate_f32 = wgate.float()
90+ 
91+ # 确定总 token 数和 batch/seq 维度（仅 BSH 布局需要保留 shape 信息用于最后 reshape）
92+ if is_th_layout:
93+ total_tokens = x.shape[0]
94+ batch_size = seq_len = None # TH 布局下不需要恢复 BSH shape
95+ else:
96+ batch_size, seq_len, hidden_size = x.shape
97+ total_tokens = batch_size * seq_len
98+ 
99+ # ══════════════════════════════════════════════════════════════════════
100+ # 阶段 B: 分配梯度累积缓冲区
101+ # ══════════════════════════════════════════════════════════════════════
102+ #
103+ # 前向计算中有三个"源头"张量会产生梯度：
104+ # (a) new_kv = x @ wkv.T — 仅仅来自 matmul
105+ # (b) new_score = x @ wgate.T + APE — 来自 matmul + 加法
106+ # (c) APE 本身 — 位置编码参数
107+ #
108+ # 在反向过程中，我们先将每个压缩块的梯度反向传播到 (a)(b)(c) 的对应
109+ # token/位置，得到 d_new_kv、d_new_score、d_ape。最后一步再将这些
110+ # 累积梯度通过 matmul 反向得到 d_x, d_wkv, d_wgate。
111+ 
112+ # 维度约定: T=total_tokens, C=coff*head_dim, H=hidden_size
113+ #
114+ # d_new_kv: (T, C) ∂Loss/∂new_kv — 仅来自 matmul 反向
115+ d_new_kv = torch.zeros(total_tokens, coff * head_dim, device=x.device)
116+ 
117+ # d_new_score: (T, C) ∂Loss/∂new_score — 同时流入 matmul 反向和 APE 反向
118+ d_new_score = torch.zeros(total_tokens, coff * head_dim, device=x.device)
119+ 
120+ # d_ape: (cmp_ratio, C) ∂Loss/∂ape
121+ d_ape = torch.zeros(cmp_ratio, coff * head_dim, device=x.device)
122+ 
123+ # ══════════════════════════════════════════════════════════════════════
124+ # 阶段 C: 逐 batch、逐压缩块进行反向传播
125+ # ══════════════════════════════════════════════════════════════════════
126+ #
127+ # 遍历顺序与前向完全一致：先遍历 batch，再在每个 batch 内按 cmp_ratio
128+ # 步长遍历压缩块。对于每个压缩块，执行反步 1→3（参见文件头部推导）。
129+ #
130+ # grad_output_index / bsh_block_idx: 按压缩块顺序消费三个输入（见 C.4a）
131+ grad_output_index = 0
132+ 
133+ for batch_idx in range(num_batches):
134+ # ── C.1 确定当前 batch 的序列范围 ────────────────────────────────
135+ # batch_start: 当前 batch 在全局序列中的绝对起始位置
136+ batch_start = start_positions[batch_idx]
137+ 
138+ # batch_seq_used: 当前 batch 实际参与计算的 token 数
139+ # seqused 非 None 时优先使用，否则用 sequence length
140+ if seqused is not None:
141+ batch_seq_used = seqused[batch_idx].item()
142+ elif is_th_layout:
143+ batch_seq_used = int(cu_seqlens[batch_idx + 1].item() - cu_seqlens[batch_idx].item())
144+ else:
145+ batch_seq_used = x.shape[1]
146+ 
147+ # compress_limit: 需要压缩的序列位置上限（向下对齐到 cmp_ratio 的整数倍）
148+ # 例如: batch_start=5, seq_used=10, cmp_ratio=4 → limit = (5+10)//4*4 = 12
149+ # 位置 5,6,7,8,9,10,11 会被压缩，12,13,14 不压缩（不足一个完整块）
150+ compress_limit = (batch_start + batch_seq_used) // cmp_ratio * cmp_ratio
151+ 
152+ # batch_seq_idx: 当前在 batch 内已处理到的相对位置（相对于 batch_start）
153+ batch_seq_idx = 0
154+ # bsh_block_idx: BSH 布局下当前 batch 的压缩块序号（从 0 开始）
155+ bsh_block_idx = 0
156+ 
157+ while batch_seq_idx < batch_seq_used:
158+ # ── C.2 当前压缩块的全局序列范围 ─────────────────────────────
159+ # block_seq_start: 当前块在全局序列中的起始位置
160+ block_seq_start = batch_start + batch_seq_idx
161+ 
162+ # block_seq_end: 当前块在全局序列中的结束位置（向前对齐到 cmp_ratio 整数倍）
163+ # 对于最后一个不完整块，end 会被截断到 batch 结尾
164+ block_seq_end = (block_seq_start // cmp_ratio) * cmp_ratio + cmp_ratio
165+ if block_seq_end > batch_start + batch_seq_used:
166+ block_seq_end = batch_start + batch_seq_used
167+ 
168+ # 推进到下一块（在 guard clause 之前更新，避免重复）
169+ batch_seq_idx = block_seq_end - batch_start
170+ 
171+ # ── C.4 尾块跳过（不足 cmp_ratio 个 token，仅写 cache）──
172+ if block_seq_start >= compress_limit:
173+ break
174+ 
175+ # ═══════════ 反向传播：反步 1→3 ═══════════
176+ 
177+ # ── 获取当前块的上游梯度和前向中间变量 ──
178+ if is_th_layout:
179+ d_compressed_kv = d_cpm_kv[grad_output_index].float()
180+ saved_kv = kv[grad_output_index].float()
181+ saved_softmax = softmax_score[grad_output_index].float()
182+ grad_output_index += 1
183+ else:
184+ d_compressed_kv = d_cpm_kv[batch_idx, bsh_block_idx].float()
185+ saved_kv = kv[batch_idx, bsh_block_idx].float()
186+ saved_softmax = softmax_score[batch_idx, bsh_block_idx].float()
187+ bsh_block_idx += 1
188+ 
189+ # ══════════════════════════════════════════════════════════
190+ # 【反步 1】element-wise mul + reduce_sum 反向
191+ # ══════════════════════════════════════════════════════════
192+ #
193+ # 前向: C = Σ_i (K_i ⊙ W_i) K=(N, hd), W=(N, hd), C=(1, hd)
194+ # 其中 N=coff*cmp_ratio, hd=head_dim
195+ # 反向: dK_i = dC ⊙ W_i, dW_i = dC ⊙ K_i (广播: dC (1,hd)→(N,hd))
196+ 
197+ d_compressed_kv_2d = d_compressed_kv.unsqueeze(0) # (1,hd) → (N,hd)
198+ d_kv_block = d_compressed_kv_2d * saved_softmax # (N,hd)
199+ d_score_weighted = d_compressed_kv_2d * saved_kv # (N,hd)
200+ # ══════════════════════════════════════════════════════════
201+ # 【反步 2】softmax 反向 (dim=0, 每列独立)
202+ # ══════════════════════════════════════════════════════════
203+ #
204+ # 前向: S = softmax(Z, dim=0) S,Z ∈ (N,hd)
205+ # 反向: dZ = S ⊙ (dS - column_sum(S ⊙ dS))
206+ # S ⊙ dS → (N,hd)
207+ # column_sum → (1,hd)
208+ # dZ → (N,hd)
209+ 
210+ softmax_backward_sum = (saved_softmax * d_score_weighted).sum(dim=0, keepdim=True) # (1,hd)
211+ d_score_block = saved_softmax * (d_score_weighted - softmax_backward_sum) # (N,hd)
212+ d_score_block = d_score_block.view(cmp_ratio, coff, head_dim)
213+ d_kv_block = d_kv_block.view(cmp_ratio, coff, head_dim)
214+ 
215+ # ══════════════════════════════════════════════════════════
216+ # 【反步 3】梯度路由 → 映射回 flat new_kv/new_score 数组
217+ # ══════════════════════════════════════════════════════════
218+ #
219+ # d_kv_block, d_score_block ∈ (N, hd) 其中 N=coff*cmp_ratio
220+ # 路由目标: d_new_kv, d_new_score ∈ (T, C) 其中 C=coff*hd
221+ #
222+ # kv_2d 的 N 行按 overlap_id 分组:
223+ # overlap_id=0: 行 [0, cmp_ratio) → 列 [0, hd)
224+ # overlap_id=coff-1: 行 [(coff-1)*cr, N) → 列 [(coff-1)*hd, C)
225+ 
226+ # flat 偏移: 当前块在展平 (T, C) 数组中的 token 行范围
227+ flat_offset_base = (
228+ int(cu_seqlens[batch_idx].item()) if is_th_layout
229+ else batch_idx * x.shape[1]
230+ )
231+ flat_offset_start = flat_offset_base + (block_seq_start - batch_start)
232+ flat_offset_end = flat_offset_base + (block_seq_end - batch_start)
233+ 
234+ # ── 3a 当前 overlap 槽位 (overlap_id = coff - 1) ──
235+ current_overlap_id = coff - 1
236+ current_dim_start = current_overlap_id * head_dim
237+ current_dim_end = (current_overlap_id + 1) * head_dim
238+ 
239+ current_tokens_from_cache = 0
240+ if batch_start == block_seq_start:
241+ current_tokens_from_cache = batch_start % cmp_ratio
242+ 
243+ block_row_start = current_overlap_id * cmp_ratio + current_tokens_from_cache
244+ block_row_end = current_overlap_id * cmp_ratio + cmp_ratio
245+ 
246+ # d_kv_block[rows, :] → d_new_kv[tokens, cols] (N,hd)子块 → (T,C)子块
247+ d_new_kv[
248+ flat_offset_start:flat_offset_end,
249+ current_dim_start:current_dim_end,
250+ ] += d_kv_block[current_tokens_from_cache:cmp_ratio, coff-1, :]
251+ 
252+ d_new_score[
253+ flat_offset_start:flat_offset_end,
254+ current_dim_start:current_dim_end,
255+ ] += d_score_block[current_tokens_from_cache:cmp_ratio, coff-1, :]
256+ 
257+ # d_ape: (cmp_ratio, C) 按 pos%cmp_ratio 累加 d_score_block 对应行
258+ ape_offset = block_seq_start % cmp_ratio
259+ n_pos = block_seq_end - block_seq_start
260+ 
261+ d_ape[ape_offset:ape_offset + n_pos, current_dim_start:current_dim_end] += d_score_block[current_tokens_from_cache:cmp_ratio, coff-1, :]
262+ 
263+ # ── 3b 前一个 overlap 槽位 (overlap_id = 0, 仅 coff == 2) ──
264+ if coff == 2:
265+ prev_tokens_from_cache = 0
266+ if batch_start == block_seq_start:
267+ prev_tokens_from_cache = cmp_ratio
268+ elif block_seq_start - cmp_ratio < batch_start:
269+ prev_tokens_from_cache = batch_start % cmp_ratio
270+ 
271+ # 全来自 cache，无需路由
272+ if prev_tokens_from_cache >= cmp_ratio:
273+ continue
274+ 
275+ prev_dim_start = 0
276+ prev_dim_end = head_dim
277+ block_row_start = prev_tokens_from_cache
278+ block_row_end = cmp_ratio
279+ 
280+ prev_flat_offset_start = flat_offset_start - (cmp_ratio - prev_tokens_from_cache)
281+ prev_flat_offset_end = flat_offset_start
282+ 
283+ # (N,hd)子块 → (T,C)子块，前一块列偏移 [0, hd)
284+ d_new_kv[
285+ prev_flat_offset_start:prev_flat_offset_end,
286+ prev_dim_start:prev_dim_end,
287+ ] += d_kv_block[prev_tokens_from_cache:cmp_ratio, 0, :]
288+ 
289+ d_new_score[
290+ prev_flat_offset_start:prev_flat_offset_end,
291+ prev_dim_start:prev_dim_end,
292+ ] += d_score_block[prev_tokens_from_cache:cmp_ratio, 0, :]
293+ 
294+ n_pos = cmp_ratio - prev_tokens_from_cache
295+ d_ape[prev_tokens_from_cache:cmp_ratio, prev_dim_start:prev_dim_end] += d_score_block[prev_tokens_from_cache:cmp_ratio, 0, :]
296+ 
297+ # ══════════════════════════════════════════════════════════════════════
298+ # 阶段 D: Matmul 反向
299+ # ══════════════════════════════════════════════════════════════════════
300+ #
301+ # 维度约定: T=total_tokens, C=coff*head_dim, H=hidden_size
302+ #
303+ # 前向有两条独立的 matmul:
304+ # new_kv = x @ wkv.T → (T, H) @ (H, C) = (T, C) 不含 APE
305+ # new_score = x @ wgate.T → (T, H) @ (H, C) = (T, C) APE 后加
306+ #
307+ # Matmul 反向公式 (Y = X @ W.T):
308+ # dX = dY @ W → (T, C) @ (C, H) = (T, H)
309+ # dW = dY.T @ X → (C, T) @ (T, H) = (C, H)
310+ #
311+ # x 同时流入两条路径，d_x 是两条之和:
312+ # d_x = d_new_kv @ wkv + d_new_score @ wgate
313+ # (T,C)@(C,H)→(T,H) (T,C)@(C,H)→(T,H)
314+ # d_wkv = d_new_kv.T @ x_flat → (C, T) @ (T, H) = (C, H)
315+ # d_wgate = d_new_score.T @ x_flat → (C, T) @ (T, H) = (C, H)
316+ 
317+ x_flat = x_f32 if is_th_layout else x_f32.reshape(total_tokens, x.shape[-1]) # (T, H)
318+ iodtype = x.dtype
319+ d_new_kv = d_new_kv.to(iodtype).float()
320+ d_new_score = d_new_score.to(iodtype).float()
321+ # print(d_new_kv.shape)
322+ # print("dkv:", d_new_kv[0])
323+ # print("d_new_kv @ wkv_f32:")
324+ # print((d_new_kv @ wkv_f32).shape, (d_new_kv @ wkv_f32)[0])
325+ # print("d_new_score @ wgate_f32:")
326+ # print((d_new_score @ wgate_f32).shape, (d_new_score @ wgate_f32)[0])
327+ d_x_flat = d_new_kv @ wkv_f32 + d_new_score @ wgate_f32 # (T, H)
328+ # print("d_x_flat:")
329+ # print(d_x_flat.shape, d_x_flat[0])
330+ d_wkv = d_new_kv.T @ x_flat # (C, H)
331+ d_wgate = d_new_score.T @ x_flat # (C, H)
332+ 
333+ # 根据原始输入布局恢复 d_x 的形状，并转换回输入的 dtype
334+ if is_th_layout:
335+ d_x = d_x_flat.to(x.dtype)
336+ else:
337+ d_x = d_x_flat.reshape(batch_size, seq_len, hidden_size).to(x.dtype)
338+ 
339+ return d_x, d_wkv.to(wkv.dtype), d_wgate.to(wgate.dtype), d_ape
Aattention/compressor_grad/tests/pytest/compressor_grad_npu.py+203-0
@@ -0,0 +1,203 @@
1+#!/usr/bin/python
2+# -*- coding: utf-8 -*-
3+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+# See LICENSE in the root of the software repository for the full text of the License.
10+# ======================================================================================================================
11+ 
12+"""NPU 后端 — compressor 反向 (compressor_grad) 自动调用测试。
13+ 
14+通过 torch.autograd 自动调用 compressor_grad:
15+1. 前向调用 compressor op，启用 requires_grad
16+2. 用上游梯度 d_cpm_kv 构造 loss
17+3. loss.backward() 自动触发 _compressor_backward → compressor_grad 算子
18+4. 收集 x / wkv / wgate / ape 的梯度输出
19+ 
20+"""
21+ 
22+import torch
23+from typing import Any, Dict, Optional
24+ 
25+try:
26+ import torch_npu # noqa: F401
27+ _HAS_TORCH_NPU = True
28+except ImportError:
29+ _HAS_TORCH_NPU = False
30+ 
31+from cann_ops_transformer.ops.compressor import _compressor_forward
32+ 
33+ 
34+def _valid_block_mask(cmp_kv, cu_seqlens, seqused, start_pos, cmp_ratio, seq_size=None):
35+ """有效压缩块 mask（与 totalValid 计算一致；padding 无效块置 False）。
36+ 
37+ 返回与 cmp_kv 同形状的 bool 张量：BSH (B, blocks, D) 每 batch 前 vb 块有效；
38+ TH (rows, D) 前 totalValid 行有效（块索引全局连续）。
39+ """
40+ d3 = cmp_kv.dim() == 3
41+ B = cmp_kv.shape[0] if d3 else (int(cu_seqlens.shape[0] - 1) if cu_seqlens is not None else 0)
42+ valid = torch.zeros_like(cmp_kv, dtype=torch.bool)
43+ if d3: # BSH: (B, blocks, D)
44+ for i in range(B):
45+ sp = int(start_pos[i]) if start_pos is not None else 0
46+ sq = int(seqused[i]) if seqused is not None else seq_size
47+ cmp_limit = (sp + sq) // cmp_ratio * cmp_ratio
48+ vb = (cmp_limit - sp + cmp_ratio - 1) // cmp_ratio if cmp_limit > sp else 0
49+ valid[i, :vb, :] = True
50+ else: # TH: (rows, D)
51+ rows = cmp_kv.shape[0]
52+ total = 0
53+ for i in range(B):
54+ sp = int(start_pos[i]) if start_pos is not None else 0
55+ sq = int(seqused[i]) if seqused is not None else \
56+ (int(cu_seqlens[i + 1] - cu_seqlens[i]) if cu_seqlens is not None else 0)
57+ cmp_limit = (sp + sq) // cmp_ratio * cmp_ratio
58+ if cmp_limit > sp:
59+ total += (cmp_limit - sp + cmp_ratio - 1) // cmp_ratio
60+ valid[:min(total, rows), :] = True
61+ return valid
62+ 
63+ 
64+class NPUBackend:
65+ """compressor_grad NPU 后端 — 通过 autograd 自动调用反向算子。
66+ 
67+ 用法:
68+ 
69+ backend = NPUBackend(device_id=0)
70+ grads = backend.compute(inputs)
71+ # grads["d_x"], grads["d_wkv"], grads["d_wgate"], grads["d_ape"]
72+ """
73+ 
74+ name = "npu"
75+ 
76+ def __init__(self, device_id: int = 0):
77+ self._device_id = device_id
78+ self._device = torch.device(f"npu:{device_id}")
79+ self._compressor_fn = None
80+ if _HAS_TORCH_NPU:
81+ torch.npu.set_device(device_id)
82+ 
83+ @property
84+ def device(self) -> torch.device:
85+ return self._device
86+ 
87+ def is_available(self) -> bool:
88+ return _HAS_TORCH_NPU
89+ 
90+ def clear_cache(self):
91+ if _HAS_TORCH_NPU:
92+ torch.npu.empty_cache()
93+ 
94+ def compute(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
95+ """通过 compressor 前向 + autograd backward 调用 compressor_grad。
96+ 
97+ 通路说明:
98+ - torch.ops.cann_ops_transformer._compressor_forward 返回
99+ (cmp_kv, softmax_score, kv) 并注册了 autograd SetupContext。
100+ - backward 中调用 _compressor_backward → csrc CompressorBwd
101+ 即 compressor_grad 算子，自动计算 d_x/d_wkv/d_wgate/d_ape。
102+ 
103+ Args:
104+ inputs:
105+ x: 输入张量 [T, H] 或 [B, S, H]
106+ wkv: KV 投影权重 (coff * D, H)
107+ wgate: Gate 投影权重 (coff * D, H)
108+ state_cache: state cache (block_num, block_size, 2*coff*D)
109+ ape: 位置编码 (cmp_ratio, coff * D)
110+ block_table: cache 模式 1 的 block table
111+ cmp_ratio: 压缩率 (默认 4)
112+ coff: 重叠因子 (默认 1)
113+ cache_mode: cache 模式 (默认 1)
114+ input_layout: "BSH" 或 "TH"
115+ cu_seqlens: TH 布局下的累积序列长度
116+ seqused: 每个 batch 的序列已用长度
117+ start_pos: 起始位置
118+ d_cpm_kv: 上游梯度 ∂Loss/∂cmp_kv — 用于 backward
119+ Returns:
120+ d_x: x 的梯度 ∂Loss/∂x
121+ d_wkv: wkv 的梯度 ∂Loss/∂wkv
122+ d_wgate: wgate 的梯度 ∂Loss/∂wgate
123+ d_ape: ape 的梯度 ∂Loss/∂ape
124+ """
125+ x = inputs["x"]
126+ wkv = inputs["wkv"]
127+ wgate = inputs["wgate"]
128+ state_cache = inputs["state_cache"]
129+ ape = inputs["ape"]
130+ d_cpm_kv = inputs["d_cpm_kv"]
131+ 
132+ cmp_ratio = inputs.get("cmp_ratio", 4)
133+ coff = inputs.get("coff", 1)
134+ cache_mode = inputs.get("cache_mode", 1)
135+ layout = inputs.get("input_layout", "BSH")
136+ 
137+ cu_seqlens = inputs.get("cu_seqlens")
138+ seqused = inputs.get("seqused")
139+ start_pos = inputs.get("start_pos")
140+ block_table = inputs.get("block_table")
141+ 
142+ x = x.to(self._device).detach()
143+ wkv = wkv.to(self._device).detach()
144+ wgate = wgate.to(self._device).detach()
145+ state_cache = state_cache.to(self._device)
146+ ape = ape.to(self._device)
147+ d_cpm_kv = d_cpm_kv.to(self._device)
148+ 
149+ cu_seqlens = cu_seqlens.to(self._device) if cu_seqlens is not None else None
150+ seqused = seqused.to(self._device) if seqused is not None else None
151+ start_pos = start_pos.to(self._device) if start_pos is not None else None
152+ block_table = block_table.to(self._device) if block_table is not None else None
153+ 
154+ # 设置 requires_grad 以启用 autograd 反向传播
155+ x.requires_grad_(True)
156+ wkv.requires_grad_(True)
157+ wgate.requires_grad_(True)
158+ ape.requires_grad_(True)
159+
160+ cmp_kv, softmax_score, kv = _compressor_forward(
161+ x, wkv, wgate, state_cache, ape,
162+ state_block_table=block_table,
163+ cu_seqlens=cu_seqlens, seqused=seqused, start_pos=start_pos,
164+ cmp_ratio=cmp_ratio, coff=coff, cache_mode=cache_mode,
165+ grad_enabled=True
166+ )
167+ # backward 会释放 autograd 保存的中间量（sm/kv）→ 提前 clone 脱离图
168+ softmax_score_saved = softmax_score.detach().clone()
169+ kv_saved = kv.detach().clone()
170+ torch.npu.synchronize()
171+ 
172+ # 用上游梯度 d_cpm_kv 构造 loss → backward 自动调用 compressor_grad
173+ # ⚠️ 只对有效压缩块计算：padding 无效块在 cmp_kv 中未写（可能垃圾/NaN），
174+ # 必须 mask 掉，否则其值污染 loss 与梯度（参考正向 cmp_mask 语义）
175+ loss = (cmp_kv * d_cpm_kv * _valid_block_mask(cmp_kv, cu_seqlens, seqused,
176+ start_pos, cmp_ratio,
177+ seq_size=x.shape[1] if x.dim() == 3 else None)).sum()
178+ loss.backward(retain_graph=True)
179+ torch.npu.synchronize()
180+ 
181+ d_x = x.grad
182+ d_wkv = wkv.grad
183+ d_wgate = wgate.grad
184+ d_ape = ape.grad
185+ 
186+ if d_x is None:
187+ d_x = torch.zeros_like(x)
188+ if d_wkv is None:
189+ d_wkv = torch.zeros_like(wkv)
190+ if d_wgate is None:
191+ d_wgate = torch.zeros_like(wgate)
192+ if d_ape is None:
193+ d_ape = torch.zeros(cmp_ratio, wkv.size(0), device=self._device, dtype=torch.float32)
194+ 
195+ return {
196+ "d_x": d_x,
197+ "d_wkv": d_wkv,
198+ "d_wgate": d_wgate,
199+ "d_ape": d_ape,
200+ "cmp_kv": cmp_kv.detach(),
201+ "softmax_score": softmax_score_saved,
202+ "kv": kv_saved,
203+ }
Aattention/compressor_grad/tests/pytest/conftest.py+17-0
@@ -0,0 +1,17 @@
1+# ======================================================================================================================
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# ======================================================================================================================
10+ 
11+ 
12+def pytest_addoption(parser):
13+ """选择验证通路：1=正反向全链路（通路5+通路4拼接，默认） 2=单反向PyPTO直跑
14+ 3=小算子拼接golden 4=单反向直调（_compressor_backward） 5=单正向直调（_compressor_forward）"""
15+ parser.addoption("--pathway", action="store", type=int, default=1,
16+ help="验证通路: 1=全链路(默认) 2=单反向PyPTO直跑 3=小算子拼接golden "
17+ "4=单反向直调 5=单正向直调")
Aattention/compressor_grad/tests/pytest/test_compressor_grad.py+985-0
@@ -0,0 +1,985 @@
1+#!/usr/bin/python
2+# -*- coding: utf-8 -*-
3+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+# See LICENSE in the root of the software repository for the full text of the License.
10+# ======================================================================================================================
11+ 
12+import random
13+ 
14+import os
15+import sys
16+import torch
17+import pytest
18+def run_small_ops_case(case, device):
19+ """通路 3：小算子拼接通路——CPU 小算子拼接 + 自动反向推导梯度作 golden，
20+ 算子侧经 pta 反向传播串联正反向（与通路 1 相同 NPU 流程）。
21+ 
22+ ⚠️ loss 只对有效压缩块计算（cmp_mask 过滤），无效 padding 不附加梯度。
23+ """
24+ p = _parse_case(case, device)
25+ B, S1, H, D = p["B"], p["S1"], p["H"], p["D"]
26+ cr, coff = p["cr"], p["coff"]
27+ layoutId, ioDtype, dataType = p["layoutId"], p["ioDtype"], p["dataType"]
28+ layout = p["layout"]
29+ seqSize = p["seqSize"]
30+ tokenSize = p["tokenSize"]
31+ sps = case["start_pos"] or [0] * B
32+ sqs = case["seqused_q"]
33+ cu = case.get("seqlens_list_q")
34+ seed = case["manual_seed"]
35+ 
36+ if layoutId == 1:
37+ if cu is not None:
38+ cuNpu = torch.tensor(cu, dtype=torch.int32)
39+ elif sqs:
40+ cuNpu = torch.tensor([0] + torch.tensor(sqs).cumsum(0).tolist(), dtype=torch.int32)
41+ else:
42+ cuNpu = torch.zeros(B + 1, dtype=torch.int32)
43+ else:
44+ cuNpu = None
45+ totalValid = 0
46+ for i in range(B):
47+ sp = sps[i]
48+ sq = sqs[i] if sqs else (int(cuNpu[i + 1] - cuNpu[i]) if cuNpu is not None
49+ else (tokenSize if layoutId == 1 else tokenSize // B))
50+ cmpLimit = (sp + sq) // cr * cr
51+ if cmpLimit > sp:
52+ totalValid += _ceil_div(cmpLimit - sp, cr)
53+ 
54+ # ── NPU 侧：pta 反向传播串联正反向（与通路 1 相同的输入构造 + compute）──
55+ b = _build_inputs(case, p, device, seed, totalValid, cuNpu, sqs, sps)
56+ backend = NPUBackend(device_id=DEVICE_ID)
57+ npu_result = backend.compute(inputs=b["inputs"])
58+ torch.npu.synchronize()
59+ 
60+ # ── CPU 侧：小算子拼接 + 自动反向推导梯度（golden）──
61+ coffD = coff * D
62+ stateCpu = b["state_cache_npu"].cpu()
63+ xCpu = b["xNpu"].cpu().view(B, seqSize, H) if layoutId == 0 else b["xNpu"].cpu()
64+ xIn = xCpu.detach().requires_grad_(True)
65+ wkvIn = b["wkvInp"].cpu().detach().requires_grad_(True)
66+ wgateIn = b["wgateInp"].cpu().detach().requires_grad_(True)
67+ apeIn = b["ape_npu"].cpu().detach().requires_grad_(True)
68+ ref_cmp, ref_mask, _, _, _, _ = compressor_forward.compressor_forward(
69+ xIn, wkvIn, wgateIn,
70+ stateCpu[..., :coffD].detach(), stateCpu[..., coffD:].detach(),
71+ apeIn,
72+ b["block_table_npu"].cpu(),
73+ sps, cmp_ratio=cr, coff=coff, cache_mode=b["cache_mode"],
74+ cu_seqlens=cuNpu.cpu() if cuNpu is not None else None,
75+ seqused=sqs, return_intermediates=True)
76+ # ⚠️ loss 只对有效压缩块：mask 过滤，无效 padding 不附加梯度
77+ dcCpu = b["dcNpu"].cpu().float()
78+ if ref_cmp.requires_grad:
79+ loss = (ref_cmp * dcCpu * ref_mask.float()).sum()
80+ loss.backward()
81+ g_dx = xIn.grad.reshape(-1, H) if layoutId == 0 else xIn.grad
82+ g_dwkv = wkvIn.grad
83+ g_dwgate = wgateIn.grad
84+ g_dape = apeIn.grad
85+ else:
86+ # 无有效压缩块（totalValid=0）：cmp_kv 为纯 zeros 无梯度图，golden 梯度全 0
87+ g_dx = torch.zeros_like(xIn).reshape(-1, H) if layoutId == 0 else torch.zeros_like(xIn)
88+ g_dwkv = torch.zeros_like(wkvIn)
89+ g_dwgate = torch.zeros_like(wgateIn)
90+ g_dape = torch.zeros_like(apeIn)
91+ 
92+ # ── 比对：NPU 四梯度 vs CPU 自动反向四梯度 ──
93+ d_xNpu = npu_result["d_x"].cpu().float().reshape(-1, H)
94+ p_pct_thd = get_pct_thd(dataType)
95+ checks = [
96+ ("d_wkv", g_dwkv.float(), npu_result["d_wkv"].cpu().float(), dataType, p_pct_thd),
97+ ("d_wgate", g_dwgate.float(), npu_result["d_wgate"].cpu().float(), dataType, p_pct_thd),
98+ ("d_ape", g_dape, npu_result["d_ape"].cpu().float(), "float32", get_pct_thd("float32")),
99+ ("d_x", g_dx.float(), d_xNpu, dataType, p_pct_thd),
100+ ]
101+ statuses = []
102+ for name, exp, act, dt, thd in checks:
103+ r = check_one_output(name, exp, act, dt, True, 1, thd)
104+ statuses.append(f"{name}={r['status']}")
105+ status = "ERROR" if any("ERROR" in s for s in statuses) else \
106+ "FAIL" if any("FAIL" in s for s in statuses) else "PASS"
107+ return {"status": status, "detail": " ".join(statuses)}
108+ 
109+ 
110+ 
111+# PyPTO kernel 直跑（通路 2）：不经 aclnn/custom 包，运行时 JIT 编译
112+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "op_kernel"))
113+from compressor_grad import compressor_grad, CompressorGradTiling # noqa: E402
114+ 
115+from test_compressor_grad_cases import REDLINE_CASES
116+from compressor_grad_cpu_golden import compressor_grad_golden
117+from compressor_grad_npu import NPUBackend
118+from cann_ops_transformer.ops.compressor import _compressor_forward, _compressor_backward
119+import compressor_forward
120+from compressor_grad_check import (
121+ check_one_output,
122+ check_result,
123+ print_log,
124+ get_pct_thd,
125+ format_case_line,
126+)
127+ 
128+DEVICE_ID = 0
129+ 
130+ 
131+def _ceil_div(a, b):
132+ return (a + b - 1) // b
133+ 
134+# ── 输出精度对比开关 ──
135+ENABLE_DWKV_CHECK = True
136+ENABLE_DWGATE_CHECK = True
137+ENABLE_APE_CHECK = True
138+ENABLE_DX_CHECK = True
139+ 
140+# ================================================================
141+# Workspace reduction
142+# ================================================================
143+ 
144+def reduce_ape_workspace(workspace_raw, headDim, coff, cmpRatio, roundCnt, totalBlocks):
145+ """
146+ Reduce d_ape workspace partials to final d_ape: shape (cmpRatio, coff*headDim).
147+ 
148+ Each Vec core writes a tile of shape [coff*cmpRatio, dBaseSize] per round,
149+ covering its D-slice columns. All slots summed produce the final d_ape.
150+ """
151+ coreNum = 64
152+ dBaseSize = 128 // coff
153+ groupSize = headDim // 128
154+ nCorePerGroup = (coreNum // 2) // groupSize
155+ mPerRoundPerGroup = 128 * 2 // coff // cmpRatio
156+ mBaseSize = 128 // cmpRatio
157+ mPerRoundAll = nCorePerGroup * mPerRoundPerGroup
158+ 
159+ dbRows = mPerRoundAll * cmpRatio
160+ ws = workspace_raw[:2 * dbRows * headDim].view(2 * dbRows, headDim).float()
161+ d_ape = torch.zeros(cmpRatio, coff * headDim)
162+ 
163+ for coreIdx in range(coreNum):
164+ cubeCoreIdx = coreIdx // 2
165+ groupIdx = cubeCoreIdx // groupSize
166+ intraGroupIdx = cubeCoreIdx % groupSize
167+ nStart = intraGroupIdx * 128 + (coreIdx % coff) * dBaseSize
168+ 
169+ for roundIdx in range(roundCnt):
170+ roundBlocks = min(mPerRoundAll,
171+ totalBlocks - roundIdx * mPerRoundAll)
172+ preDealTcSize = groupIdx * mPerRoundPerGroup
173+ if coff == 1 and coreIdx % 2 == 1:
174+ preDealTcSize += mBaseSize
175+ if preDealTcSize >= roundBlocks:
176+ continue
177+ dealTcSize = min(mBaseSize,
178+ max(0, roundBlocks - preDealTcSize))
179+ if dealTcSize <= 0:
180+ continue
181+ 
182+ dbBase = (roundIdx % 2) * dbRows
183+ localSlot = groupIdx * mPerRoundPerGroup * cmpRatio
184+ if coff == 1 and coreIdx % 2 == 1:
185+ localSlot += mBaseSize * cmpRatio
186+ 
187+ # [coff*cmpRatio, dBaseSize] at (dbBase + localSlot, nStart)
188+ tile = ws[dbBase + localSlot:dbBase + localSlot + coff * cmpRatio,
189+ nStart:nStart + dBaseSize]
190+ 
191+ for oid in range(coff):
192+ row_start = oid * cmpRatio
193+ col_start = oid * headDim + nStart
194+ d_ape[:, col_start:col_start + dBaseSize] += \
195+ tile[row_start:row_start + cmpRatio, :]
196+ 
197+ # Note: with 2x db, rounds N and N+2 share the same slot.
198+ # Phase 3 will consume per-round; standalone test sees only last 2 rounds.
199+ 
200+ return d_ape
201+ 
202+# ================================================================
203+# Test runner
204+# ================================================================
205+ 
206+def _check_block_size(case, batchSize, seqSize, layoutId, cmpRatio, coff, cacheMode,
207+ cuNpu, sqs):
208+ """校验 case 数据的 block_size 字段合法性；不符则终止+告警（不静默替换）。
209+ 
210+ cache_mode=1: block_size 为自由参数，需在 [1, 1024]；
211+ cache_mode=2: block_size 必须等于 min(coff*cmpRatio + S_for_bs - 1, 1024)
212+ （S_for_bs: TH 布局 = max(seqused) 或 max(cu_seqlens 差)；BSH 布局 = S1）
213+ """
214+ bs = case["block_size"]
215+ name = case["testcase_name"]
216+ if cacheMode == 1:
217+ if not (1 <= bs <= 1024):
218+ raise SystemExit(
219+ f"[block_size 校验失败] {name}: cache_mode=1 block_size={bs} "
220+ f"不在 [1,1024] 范围内，请修正 case 数据后重跑")
221+ return
222+ # cache_mode=2
223+ if layoutId == 1: # TH
224+ if sqs:
225+ S_for_bs = max(sqs)
226+ elif cuNpu is not None:
227+ S_for_bs = max(int(cuNpu[i + 1]) - int(cuNpu[i])
228+ for i in range(batchSize)) if batchSize > 0 else 0
229+ else:
230+ S_for_bs = seqSize
231+ else: # BSH
232+ S_for_bs = seqSize
233+ need = coff * cmpRatio + S_for_bs - 1
234+ expect = min(need, 1024)
235+ if bs != expect:
236+ raise SystemExit(
237+ f"[block_size 校验失败] {name}: cache_mode=2 block_size={bs} "
238+ f"与要求 {expect} 不一致 (coff*cmpRatio+S-1={need}, S_for_bs={S_for_bs})，"
239+ f"请修正 case 数据后重跑")
240+ 
241+def check_forward_outputs(npu_result, xCpu, wkvCpu, wgateCpu, stateCacheCpu, apeCpu,
242+ blockTableCpu, spGolden, sqGolden, cuGolden,
243+ cmpRatio, coff, cacheMode, headDim, dataType, ref=None):
244+ """通路 1/5：正向输出（cmp_kv/softmax_score/kv）与小算子拼接参考比对。
245+ 
246+ 参考实现：compressor_forward.compressor_forward（纯 PyTorch 小算子拼接）。
247+ 仅比对有效压缩块行（padding 行参考为 0、NPU 不保证写）。
248+ ref: 可选预计算参考（return_intermediates=True 的 6/8 元组），避免重复计算；
249+ 为 None 时内部调用 compressor_forward。
250+ """
251+ coffD = coff * headDim
252+ kv_state = stateCacheCpu[..., :coffD]
253+ score_state = stateCacheCpu[..., coffD:]
254+ if ref is None:
255+ ref = compressor_forward.compressor_forward(
256+ xCpu, wkvCpu, wgateCpu, kv_state, score_state, apeCpu, blockTableCpu,
257+ [int(v) for v in spGolden] if spGolden is not None else None,
258+ cmp_ratio=cmpRatio, coff=coff, cache_mode=cacheMode,
259+ cu_seqlens=cuGolden,
260+ seqused=[int(v) for v in sqGolden] if sqGolden is not None else None,
261+ return_intermediates=True)
262+ ref_cmp, ref_mask, _, _, ref_sm, ref_kv = ref[:6]
263+ 
264+ valid = ref_mask.any(dim=-1) # (B, blocks) 或 (rows,)
265+ 
266+ def interleave(t):
267+ # 参考为顺序布局 (prev|cur)；NPU 正向输出为交错布局
268+ # (偶行=prev、奇行=cur)——比对前重排参考。
269+ # coff=1 时无 prev/cur 之分，参考即顺序布局，直接返回（对半分会导致
270+ # 偶数错位 / 奇数崩溃）。
271+ if coff == 1:
272+ return t
273+ n = t.shape[-2]
274+ half = n // 2
275+ return torch.stack([t[..., :half, :], t[..., half:, :]], dim=-2).reshape(t.shape)
276+ 
277+ def gather(t):
278+ t = t.detach().float() # npu 侧输出/参考输入可能带 requires_grad（autograd 图）
279+ return t[valid]
280+ 
281+ results = []
282+ per = {}
283+ fwd_key_map = {
284+ "cmp_kv": ("cmpkvPct", "cmpkvStatus"),
285+ "softmax_score": ("smPct", "smStatus"),
286+ "kv": ("kvPct", "kvStatus"),
287+ }
288+ for name, ref_t, npu_t, dt in (
289+ ("cmp_kv", ref_cmp, npu_result["cmp_kv"], dataType),
290+ ("softmax_score", interleave(ref_sm), npu_result["softmax_score"], "float32"),
291+ ("kv", interleave(ref_kv), npu_result["kv"], "float32"),
292+ ):
293+ # 与反向 4 输出同一检查函数（check_one_output）：
294+ # 输出块格式（check <名> / Loop 明细 / Rtol-Atol 汇总 / 错误值打印）完全一致
295+ r = check_one_output(name, gather(ref_t), gather(npu_t), dt, True, 1,
296+ get_pct_thd(dt))
297+ r["name"] = name
298+ results.append(r)
299+ pk, sk = fwd_key_map[name]
300+ per[pk] = r["pct"]
301+ per[sk] = r["status"]
302+ detail = " ".join(f"{r['name']}={r['pct']:.2f}%[{r['status']}]" for r in results)
303+ ok = all(r["status"] == "PASS" for r in results)
304+ d = {"status": "PASS" if ok else "FAIL", "detail": detail}
305+ d.update(per)
306+ return d
307+ 
308+ 
309+def check_state_outputs(npu_state_cache, ref_kv_state, ref_score_state, kv_mask, score_mask,
310+ coff, headDim, dataType, totalValid):
311+ """通路 5：state_cache 4 项校验（与正向 pytest 的 kv_state_update / score_state_update /
312+ kv_state_origin / score_state_origin 完全一致）。
313+ 
314+ npu_state_cache: (block_num, block_size, 2*coff*headDim) float32——直调正向后已原地更新；
315+ 前一半列 = kv_state（stateIdx=0），后一半列 = score_state（stateIdx=1）。
316+ ref_kv_state / ref_score_state: 参考（compressor_forward 写回后）的 CPU state。
317+ kv_mask / score_mask: 参考的写入位置 bool mask（与正向 pytest 的 update_kv/update_score 同语义）。
318+ 
319+ 校验：
320+ *update: 写入位置 NPU vs 参考（pct_thd 常规阈值）；
321+ *origin: 未写位置 NPU vs 参考（= 初始值），pct_thd=0.0 严格相等（kernel 不得误写）。
322+ """
323+ coffD = coff * headDim
324+ kv_half = npu_state_cache[..., :coffD]
325+ score_half = npu_state_cache[..., coffD:]
326+ pct_thd = get_pct_thd("float32")
327+ results = []
328+ per = {}
329+ state_key_map = {
330+ "kv_state_update": ("kvUpdPct", "kvUpdStatus"),
331+ "score_state_update": ("scoreUpdPct", "scoreUpdStatus"),
332+ "kv_state_origin": ("kvOrgPct", "kvOrgStatus"),
333+ "score_state_origin": ("scoreOrgPct", "scoreOrgStatus"),
334+ }
335+ for name, exp, act, thd in (
336+ ("kv_state_update", ref_kv_state[kv_mask], kv_half[kv_mask], pct_thd),
337+ ("score_state_update", ref_score_state[score_mask], score_half[score_mask], pct_thd),
338+ ("kv_state_origin", ref_kv_state[~kv_mask], kv_half[~kv_mask], 0.0),
339+ ("score_state_origin", ref_score_state[~score_mask], score_half[~score_mask], 0.0),
340+ ):
341+ r = check_one_output(name, exp, act, dataType, True, totalValid, thd)
342+ results.append((name, r))
343+ pk, sk = state_key_map[name]
344+ per[pk] = r["pct"]
345+ per[sk] = r["status"]
346+ detail = " ".join(f"{n}={r['pct']:.2f}%[{r['status']}]" for n, r in results)
347+ ok = all(r["status"] == "PASS" for _, r in results)
348+ d = {"status": "PASS" if ok else "FAIL", "detail": detail}
349+ d.update(per)
350+ return d
351+ 
352+def _prepare_case(case, device):
353+ """通路 1/4/5 共用：解析 case + 构造 cuNpu/totalValid/sps/sqs/seed。"""
354+ p = _parse_case(case, device)
355+ batchSize, layoutId = p["B"], p["layoutId"]
356+ tokenSize = p["tokenSize"]
357+ sps = case["start_pos"]
358+ if sps is None:
359+ sps = [0] * batchSize
360+ sqs = case["seqused_q"]
361+ seqlens_list_q = case["seqlens_list_q"]
362+ seed = case["manual_seed"]
363+ 
364+ # ── 构造 cu_seqlens（TH 布局；BSH 为 None）──
365+ if layoutId == 1:
366+ # case 数据统一为"0 开始累加"的累积形式（B+1 项、首项 0），直用不再调整
367+ if seqlens_list_q is not None:
368+ cuNpu = torch.tensor(seqlens_list_q, dtype=torch.int32)
369+ elif sqs is not None:
370+ cuNpu = torch.tensor([0] + torch.tensor(sqs).cumsum(0).tolist(),
371+ dtype=torch.int32)
372+ else:
373+ cuNpu = torch.zeros(batchSize + 1, dtype=torch.int32)
374+ else:
375+ cuNpu = None
376+ 
377+ # ── Total valid compression blocks ──
378+ totalValid = 0
379+ for i in range(batchSize):
380+ sp = sps[i] if sps else 0
381+ if sqs:
382+ sq = sqs[i]
383+ elif cuNpu is not None:
384+ sq = int(cuNpu[i + 1].item() - cuNpu[i].item())
385+ else:
386+ sq = tokenSize if layoutId == 1 else tokenSize // batchSize
387+ cmpLimit = (sp + sq) // p["cr"] * p["cr"]
388+ if cmpLimit > sp:
389+ # startPos%cmpRatio!=0 时首块为部分块但仍计 1 块 → 向上取整（与 kernel 一致）
390+ totalValid += (cmpLimit - sp + p["cr"] - 1) // p["cr"]
391+ return p, totalValid, cuNpu, sps, sqs, seed
392+ 
393+ 
394+def run_forward_direct_case(case, p, b, totalValid, device):
395+ """通路 5：单正向——直调 _compressor_forward（不经 autograd）。
396+ 
397+ 校验：
398+ - 正向 3 输出：cmp_kv / softmax_score / kv vs compressor_forward 参考（check_forward_outputs）；
399+ - state_cache 4 项：kv_state_update / score_state_update / kv_state_origin /
400+ score_state_origin（check_state_outputs，与正向 pytest 一致；origin 严格相等）。
401+ 
402+ 返回 (result, intermediates)：intermediates = {cmp_kv, softmax_score, kv} 为
403+ NPU 真实中间量，供通路 4（直调反向）复用——通路 1 = 通路 5 + 通路 4 拼接。
404+ """
405+ batchSize, seqSize, headDim, hiddenSize = p["B"], p["S1"], p["D"], p["H"]
406+ cmpRatio, coff = p["cr"], p["coff"]
407+ layoutId, dataType = p["layoutId"], p["dataType"]
408+ wkvInp, wgateInp = b["wkvInp"], b["wgateInp"]
409+ state_cache_npu = b["state_cache_npu"]
410+ ape_npu = b["ape_npu"]
411+ block_table_npu = b["block_table_npu"]
412+ cache_mode = b["cache_mode"]
413+ cuArg = b["cuArg"]
414+ if cuArg is not None:
415+ cuArg = cuArg.to(device) # _build_inputs 中 cuNpu 在 CPU 构造（golden 用）；直调必须 NPU
416+ sqNpu = b["inputs"]["seqused"]
417+ spNpu = b["inputs"]["start_pos"]
418+ sps = b["inputs"]["start_pos"]
419+ sqs = b["inputs"]["seqused"]
420+ 
421+ # 参考输入：初始 state 必须在直调 forward 之前 clone（forward 原地更新 state_cache）
422+ state_cpu_initial = state_cache_npu.cpu()
423+ 
424+ # ── 直调正向（不经 autograd）；grad_enabled=True 才会写 softmax_score/kv 中间输出 ──
425+ cmp_kv, softmax_score, kv = _compressor_forward(
426+ b["inputs"]["x"], wkvInp, wgateInp, state_cache_npu, ape_npu,
427+ state_block_table=block_table_npu,
428+ cu_seqlens=cuArg, seqused=sqNpu, start_pos=spNpu,
429+ cmp_ratio=cmpRatio, coff=coff, cache_mode=cache_mode,
430+ grad_enabled=True)
431+ torch.npu.synchronize()
432+ npu_state_updated = state_cache_npu.cpu()
433+ 
434+ # ── CPU 参考（小算子拼接 + 写回 state + update mask）──
435+ coffD = coff * headDim
436+ xCpu = b["xNpu"]
437+ xCpu = xCpu.cpu().view(batchSize, seqSize, hiddenSize) if layoutId == 0 else xCpu.cpu()
438+ cuGolden = cuArg.cpu() if cuArg is not None else None
439+ ref = compressor_forward.compressor_forward(
440+ xCpu, wkvInp.cpu(), wgateInp.cpu(),
441+ state_cpu_initial[..., :coffD], state_cpu_initial[..., coffD:],
442+ ape_npu.cpu(), block_table_npu.cpu(),
443+ [int(v) for v in sps] if sps is not None else None,
444+ cmp_ratio=cmpRatio, coff=coff, cache_mode=cache_mode,
445+ cu_seqlens=cuGolden,
446+ seqused=[int(v) for v in sqs] if sqs is not None else None,
447+ return_intermediates=True, return_update_mask=True)
448+ # ref = (cmp_kv, cmp_mask, kv_state_out, score_state_out, sm, kv, kv_mask, score_mask)
449+ 
450+ # ── 校验 ──
451+ npu_result = {"cmp_kv": cmp_kv, "softmax_score": softmax_score, "kv": kv}
452+ fwd = check_forward_outputs(npu_result, xCpu, wkvInp.cpu(), wgateInp.cpu(),
453+ state_cpu_initial, ape_npu.cpu(), block_table_npu.cpu(),
454+ sps, sqs, cuGolden,
455+ cmpRatio, coff, cache_mode, headDim, dataType, ref=ref)
456+ state = check_state_outputs(npu_state_updated, ref[2], ref[3], ref[6], ref[7],
457+ coff, headDim, dataType, totalValid)
458+ print_log(f"Forward check: {fwd['detail']}")
459+ print_log(f"State check: {state['detail']}")
460+ 
461+ result = {"fwdStatus": fwd["status"], "fwdDetail": fwd["detail"],
462+ "stateStatus": state["status"], "stateDetail": state["detail"]}
463+ for k in ("cmpkvPct", "cmpkvStatus", "smPct", "smStatus", "kvPct", "kvStatus"):
464+ result[k] = fwd[k]
465+ for k in ("kvUpdPct", "kvUpdStatus", "scoreUpdPct", "scoreUpdStatus",
466+ "kvOrgPct", "kvOrgStatus", "scoreOrgPct", "scoreOrgStatus"):
467+ result[k] = state[k]
468+ statuses = [fwd["status"], state["status"]]
469+ result["status"] = "ERROR" if any(s == "ERROR" for s in statuses) else \
470+ "FAIL" if any(s == "FAIL" for s in statuses) else \
471+ "SKIP" if all(s == "SKIP" for s in statuses) else "PASS"
472+ return result, {"cmp_kv": cmp_kv, "softmax_score": softmax_score, "kv": kv}
473+ 
474+ 
475+def run_backward_direct_case(case, p, b, totalValid, device, fwd=None):
476+ """通路 4：单反向——直调 _compressor_backward（不经 autograd）。
477+ 
478+ 中间量（softmax_score/kv）默认来自真实正向：fwd 为 None 时先直调一次
479+ _compressor_forward（与 golden 一致性原则相同——golden 输入用前向真实输出）；
480+ 通路 1 拼接时由通路 5 传入复用同一份中间量。
481+ 
482+ 校验反向 4 输出：d_wkv / d_wgate / d_ape / d_x vs compressor_grad_golden。
483+ """
484+ batchSize, seqSize, headDim, hiddenSize = p["B"], p["S1"], p["D"], p["H"]
485+ cmpRatio, coff = p["cr"], p["coff"]
486+ layoutId, dataType = p["layoutId"], p["dataType"]
487+ wkvInp, wgateInp = b["wkvInp"], b["wgateInp"]
488+ dcNpu = b["dcNpu"]
489+ inputs = b["inputs"]
490+ cuArg, sqNpu, spNpu = inputs["cu_seqlens"], inputs["seqused"], inputs["start_pos"]
491+ if cuArg is not None:
492+ cuArg = cuArg.to(device) # _build_inputs 中 cuNpu 在 CPU 构造（golden 用）；直调必须 NPU
493+ 
494+ if fwd is None:
495+ # 先直调一次真实正向拿中间量（不经 autograd）
496+ cmp_kv, softmax_score, kv = _compressor_forward(
497+ inputs["x"], wkvInp, wgateInp, b["state_cache_npu"], b["ape_npu"],
498+ state_block_table=b["block_table_npu"],
499+ cu_seqlens=cuArg, seqused=sqNpu, start_pos=spNpu,
500+ cmp_ratio=cmpRatio, coff=coff, cache_mode=b["cache_mode"],
501+ grad_enabled=True)
502+ torch.npu.synchronize()
503+ else:
504+ cmp_kv, softmax_score, kv = fwd["cmp_kv"], fwd["softmax_score"], fwd["kv"]
505+ 
506+ # ── 直调反向（不经 autograd）──
507+ d_x, d_wkv, d_wgate, d_ape = _compressor_backward(
508+ dcNpu, inputs["x"], wkvInp, wgateInp, softmax_score, kv,
509+ cu_seqlens=cuArg, seqused=sqNpu, start_pos=spNpu,
510+ cmp_ratio=cmpRatio, coff=coff)
511+ torch.npu.synchronize()
512+ 
513+ # ── CPU Golden（与通路 1 相同：golden 输入用前向真实输出）──
514+ dcCpu = dcNpu.cpu().float()
515+ spGolden = spNpu.cpu() if spNpu is not None else None
516+ sqGolden = sqNpu.cpu() if sqNpu is not None else None
517+ cuGolden = cuArg.cpu() if cuArg is not None else None
518+ xCpu = inputs["x"].cpu()
519+ if layoutId == 0: # BSH
520+ xCpu = xCpu.view(batchSize, seqSize, hiddenSize)
521+ dXGold, dWkvGold, dWgateGold, apeGold = compressor_grad_golden(
522+ x=xCpu, wkv=wkvInp.cpu().float(), wgate=wgateInp.cpu().float(),
523+ d_cpm_kv=(dcCpu.view(batchSize, -1, headDim) if layoutId == 0 else dcCpu),
524+ softmax_score=softmax_score.cpu(),
525+ kv=kv.cpu(),
526+ cu_seqlens=cuGolden, seqused=sqGolden, start_pos=spGolden,
527+ cmp_ratio=cmpRatio, coff=coff,
528+ )
529+ 
530+ # ── 反向 4 输出比对 ──
531+ p_pct_thd = get_pct_thd(dataType)
532+ dWkvNp = d_wkv.cpu().float()
533+ dWgateNp = d_wgate.cpu().float()
534+ apeNp = d_ape.cpu().view(cmpRatio, coff * headDim)
535+ dXNp = d_x.cpu().float()
536+ dXGoldF = dXGold.float()
537+ 
538+ print_log("=" * 80)
539+ print_log(f"Start precision check for case '{case['testcase_name']}' (validBlocks={totalValid})")
540+ 
541+ r_dwkv = check_one_output("d_wkv", dWkvGold.float(), dWkvNp, dataType,
542+ ENABLE_DWKV_CHECK, totalValid, p_pct_thd)
543+ r_dwgate = check_one_output("d_wgate", dWgateGold.float(), dWgateNp, dataType,
544+ ENABLE_DWGATE_CHECK, totalValid, p_pct_thd)
545+ r_ape = check_one_output("d_ape", apeGold, apeNp, "float32",
546+ ENABLE_APE_CHECK, totalValid, get_pct_thd("float32"))
547+ r_dx = check_one_output("d_x", dXGoldF, dXNp, dataType,
548+ ENABLE_DX_CHECK, totalValid, p_pct_thd)
549+ 
550+ result = dict(
551+ dwkvDiff=r_dwkv["diff"], dwkvPct=r_dwkv["pct"], dwkvStatus=r_dwkv["status"],
552+ dwgateDiff=r_dwgate["diff"], dwgatePct=r_dwgate["pct"], dwgateStatus=r_dwgate["status"],
553+ apeDiff=r_ape["diff"], apePct=r_ape["pct"], apeStatus=r_ape["status"],
554+ dxDiff=r_dx["diff"], dxPct=r_dx["pct"], dxStatus=r_dx["status"],
555+ )
556+ statuses = [r_dwkv["status"], r_dwgate["status"], r_ape["status"], r_dx["status"]]
557+ result["status"] = "ERROR" if any(s == "ERROR" for s in statuses) else \
558+ "FAIL" if any(s == "FAIL" for s in statuses) else \
559+ "SKIP" if all(s == "SKIP" for s in statuses) else "PASS"
560+ return result
561+ 
562+ 
563+def _parse_case(case, device):
564+ """解析 case 公共参数（通路 1/2/4/5 共用）。"""
565+ B, S1, H, D = case["B"], case["S1"], case["hidden_size"], case["D"]
566+ cr, coff = case["cmp_ratio"], case["coff"]
567+ layout = case["input_layout"]
568+ layoutId = 1 if layout == "TH" else 0
569+ dataType = case["dtype"]
570+ ioDtype = torch.float16 if dataType == "float16" else torch.bfloat16
571+ if layoutId == 1:
572+ tokenSize = int(case["seqlens_list_q"][-1])
573+ outputRows = min(tokenSize, tokenSize // cr + B)
574+ seqSize = 0 # 与 host tiling 一致（TH 不设置 seqSize）
575+ else:
576+ tokenSize = B * S1
577+ outputRows = B * _ceil_div(S1, cr)
578+ seqSize = S1
579+ return dict(B=B, S1=S1, H=H, D=D, cr=cr, coff=coff, layout=layout,
580+ layoutId=layoutId, dataType=dataType, ioDtype=ioDtype,
581+ tokenSize=tokenSize, seqSize=seqSize, outputRows=outputRows)
582+ 
583+ 
584+def _build_inputs(case, p, device, seed, totalValid, cuNpu, sqs, sps, state_init="zeros"):
585+ """构造通路 1/3/4/5 共用的 NPU 输入（x/wkv/wgate/state_cache/ape/dc/block_table + 输出占位）。
586+ 
587+ 返回 compute inputs dict + 附带 tensor（输出占位用 rand 模拟未初始化内存，
588+ 不依赖初始 0——kernel 未写输出时比对必须 FAIL 而非被零掩盖）。
589+ state_init: "zeros"（通路 1/3/4，反向不依赖 state 内容）或
590+ "rand"（通路 5 state 校验：与正向 pytest 的 kv_state/score_state
591+ 随机 uniform(-10,10) 初始一致，origin 校验未写位置保持原值）。
592+ """
593+ batchSize, seqSize, headDim, hiddenSize = p["B"], p["S1"], p["D"], p["H"]
594+ cmpRatio, coff = p["cr"], p["coff"]
595+ layoutId, ioDtype = p["layoutId"], p["ioDtype"]
596+ layout = p["layout"]
597+ tokenSize, outputRows = p["tokenSize"], p["outputRows"]
598+ sqNpu = torch.tensor(sqs, device=device, dtype=torch.int32) if sqs else None
599+ spNpu = torch.tensor(sps, device=device, dtype=torch.int32) if sps else None
600+ 
601+ # ── Workspace sizing（仅注释参考，autograd 路径由 aclnn 分配）──
602+ scKvRowCount = coff * cmpRatio
603+ coreNum = 64
604+ groupSize = headDim // 128
605+ nCorePerGroup = (coreNum // 2) // groupSize
606+ mPerRoundPerGroup = 128 * 2 // coff // cmpRatio
607+ mPerRoundAll = nCorePerGroup * mPerRoundPerGroup
608+ dbRows = mPerRoundAll * cmpRatio
609+ apeWsSize = 2 * dbRows * headDim
610+ dxPerGroup = 256 * hiddenSize
611+ dxPerRoundAll = 32 * dxPerGroup
612+ dxWsSize = 2 * dxPerRoundAll * 2
613+ cvGmPerCore = 128 * coff * (128 // coff)
614+ cvGmSize = 32 * cvGmPerCore * 2
615+ wsSize = apeWsSize + dxWsSize + cvGmSize
616+ 
617+ # ── Generate random inputs ──
618+ torch.manual_seed(seed)
619+ if layoutId == 0: # BSH: dc 与 cmp_kv 同为 (B, blocks, D)，kernel 按扁平读兼容
620+ dcNpu = torch.randn(batchSize, (seqSize + cmpRatio - 1) // cmpRatio, headDim,
621+ device=device, dtype=ioDtype)
622+ else:
623+ dcNpu = torch.randn(outputRows, headDim, device=device, dtype=ioDtype)
624+ 
625+ # ── Output tensors（占位 rand）──
626+ totalTokens = tokenSize if layoutId == 1 else batchSize * seqSize
627+ d_xNpu = torch.rand(totalTokens, hiddenSize, device=device, dtype=ioDtype)
628+ d_wkvNpu = torch.rand(coff * headDim, hiddenSize, device=device, dtype=torch.float32)
629+ d_wgateNpu = torch.rand(coff * headDim, hiddenSize, device=device, dtype=torch.float32)
630+ dapeWsNpu = torch.rand(wsSize, device=device, dtype=torch.float32)
631+ dapeOutNpu = torch.rand(coff * cmpRatio * headDim, device=device, dtype=torch.float32)
632+ 
633+ # cu_seqlens argument (None if BSH or noCu)
634+ cuArg = None if (layoutId == 0 or case.get('noCu')) else cuNpu
635+ 
636+ # x / wkv / wgate input tensors
637+ xNpu = torch.randn(totalTokens, hiddenSize, device=device, dtype=ioDtype)
638+ wkvInp = torch.randn(coff * headDim, hiddenSize, device=device, dtype=ioDtype)
639+ wgateInp = torch.randn(coff * headDim, hiddenSize, device=device, dtype=ioDtype)
640+ 
641+ # ── state_cache 与 block_table（构造与正向 golden 一致）──
642+ cache_mode = case.get("cache_mode", 1)
643+ if cache_mode == 1:
644+ block_size = case.get("block_size", 128)
645+ if layoutId == 1: # TH
646+ cu_lens = cuNpu.tolist() if cuNpu is not None else [0] * (batchSize + 1)
647+ if sps:
648+ S_max = max(sps[i] + (cu_lens[i + 1] - cu_lens[i])
649+ for i in range(batchSize)) if batchSize > 0 else 0
650+ else:
651+ S_max = max((cu_lens[i + 1] - cu_lens[i])
652+ for i in range(batchSize)) if batchSize > 0 else 0
653+ else: # BSH
654+ S_max = max(sps) + seqSize if sps else seqSize
655+ max_block_num_per_batch = (S_max + block_size - 1) // block_size
656+ block_num = batchSize * max_block_num_per_batch
657+ state_shape = (block_num, block_size, 2 * coff * headDim)
658+ else:
659+ # cache_mode=2: block_size 直用 case 数据（合法性由 _check_block_size 校验）
660+ state_shape = (batchSize, case["block_size"], 2 * coff * headDim)
661+ 
662+ _check_block_size(case, batchSize, seqSize, layoutId, cmpRatio, coff, cache_mode,
663+ cuNpu, sqs)
664+ 
665+ if state_init == "rand":
666+ # 与正向 pytest 的 kv_state/score_state 初始一致：uniform(-10, 10)
667+ state_cache_npu = torch.empty(state_shape, device=device, dtype=torch.float32).uniform_(-10, 10)
668+ else:
669+ state_cache_npu = torch.zeros(state_shape, device=device, dtype=torch.float32)
670+ ape_npu = torch.randn(torch.Size([cmpRatio, coff * headDim]), device=device,
671+ dtype=torch.float32)
672+ 
673+ block_table_npu = None
674+ if cache_mode == 1:
675+ block_table = torch.zeros(batchSize, max_block_num_per_batch, dtype=torch.int32)
676+ if batchSize > 0 and totalValid > 0:
677+ next_block_id = 1
678+ for i in range(batchSize):
679+ sp = sps[i] if sps else 0
680+ if sqs:
681+ sq = sqs[i]
682+ elif cuNpu is not None:
683+ sq = int(cuNpu[i + 1].item() - cuNpu[i].item())
684+ else:
685+ sq = tokenSize if layoutId == 1 else tokenSize // batchSize
686+ end_pos = sq
687+ # 读取范围 (与 golden compressor_golden.py:572-583 一致)
688+ cur_start = sp // cmpRatio * cmpRatio - cmpRatio
689+ cur_end = sp // cmpRatio * cmpRatio + cmpRatio
690+ if sp % cmpRatio == 0:
691+ cur_end = sp
692+ cur_end = min(cur_end, sp + sq)
693+ for j in range(max(cur_start // block_size, 0),
694+ (cur_end - 1) // block_size + 1):
695+ if next_block_id < block_num:
696+ block_table[i][j] = next_block_id
697+ next_block_id = next_block_id + 1
698+ # 写入范围 (与 golden compressor_golden.py:589-600 一致)
699+ next_start = (sp + end_pos) // cmpRatio * cmpRatio - cmpRatio
700+ next_end = (sp + end_pos) // cmpRatio * cmpRatio + cmpRatio
701+ if (sp + end_pos) % cmpRatio == 0:
702+ next_end = sp + end_pos
703+ next_end = min(next_end, sp + end_pos)
704+ for j in range(max(next_start // block_size, 0),
705+ (next_end - 1) // block_size + 1):
706+ if next_block_id < block_num and block_table[i][j] == 0:
707+ block_table[i][j] = next_block_id
708+ next_block_id = next_block_id + 1
709+ block_table_npu = block_table.to(device)
710+ else:
711+ block_table_npu = torch.tensor(
712+ random.sample(list(range(batchSize)), batchSize),
713+ device=device, dtype=torch.int32) if batchSize > 0 else torch.zeros(0, device=device, dtype=torch.int32)
714+ 
715+ x_npu_input = xNpu.view(batchSize, seqSize, hiddenSize) if layoutId == 0 else xNpu
716+ inputs = {
717+ "x": x_npu_input, "wkv": wkvInp, "wgate": wgateInp,
718+ "state_cache": state_cache_npu, "ape": ape_npu, "d_cpm_kv": dcNpu,
719+ "cu_seqlens": cuArg, "seqused": sqNpu, "start_pos": spNpu,
720+ "block_table": block_table_npu, "cmp_ratio": cmpRatio, "coff": coff,
721+ "cache_mode": cache_mode, "input_layout": layout,
722+ }
723+ return dict(inputs=inputs, xNpu=xNpu, wkvInp=wkvInp, wgateInp=wgateInp,
724+ dcNpu=dcNpu, state_cache_npu=state_cache_npu, ape_npu=ape_npu,
725+ block_table_npu=block_table_npu, cuArg=cuArg,
726+ d_xNpu=d_xNpu, d_wkvNpu=d_wkvNpu, d_wgateNpu=d_wgateNpu,
727+ dapeOutNpu=dapeOutNpu, totalTokens=totalTokens,
728+ cache_mode=cache_mode, scKvRowCount=scKvRowCount)
729+ 
730+ 
731+def run_case(case, device):
732+ """通路 1：正反向全链路 = 通路 5（直调正向）+ 通路 4（直调反向）拼接。
733+ 
734+ 共享同一份输入构造（_build_inputs, state 随机初始）与 NPU 真实中间量
735+ （softmax_score/kv），校验正向 3 输出 + state_cache 4 项 + 反向 4 输出。
736+ """
737+ name = case["testcase_name"]
738+ p, totalValid, cuNpu, sps, sqs, seed = _prepare_case(case, device)
739+ b = _build_inputs(case, p, device, seed, totalValid, cuNpu, sqs, sps,
740+ state_init="rand")
741+ 
742+ # 通路 5：直调正向 + 校验（cmp_kv/sm/kv + state 4 项），返回 NPU 中间量
743+ fwd_result, fwd_inter = run_forward_direct_case(case, p, b, totalValid, device)
744+ # 通路 4：直调反向（复用同一份中间量）+ 校验（反向 4 输出）
745+ bwd_result = run_backward_direct_case(case, p, b, totalValid, device, fwd=fwd_inter)
746+ 
747+ result = dict(name=name, totalValid=totalValid,
748+ coff=p["coff"], cmpRatio=p["cr"], headDim=p["D"],
749+ batchSize=p["B"])
750+ result.update(fwd_result)
751+ result.update(bwd_result)
752+ 
753+ # ── Overall status ──
754+ # 任一 FAIL → FAIL；任一 ERROR → ERROR
755+ # 注：totalValid=0 不再 trivial PASS——四输出必须真实写 0（golden 全 0），
756+ # kernel 未写输出（垃圾）时校验必须 FAIL（见 _zero_outputs）
757+ statuses = [result["dwkvStatus"], result["dwgateStatus"],
758+ result["apeStatus"], result["dxStatus"],
759+ result["fwdStatus"], result["stateStatus"]]
760+ if any(s == "ERROR" for s in statuses):
761+ result["status"] = "ERROR"
762+ elif any(s == "FAIL" for s in statuses):
763+ result["status"] = "FAIL"
764+ elif all(s == "SKIP" for s in statuses):
765+ result["status"] = "SKIP"
766+ else:
767+ result["status"] = "PASS"
768+ 
769+ return result
770+ 
771+ 
772+def run_backward_case(case, device):
773+ """通路 2：单反向——随机构造中间量 → PyPTO 直跑 kernel → golden 比对。
774+ 
775+ 不经 aclnn/custom 包（kernel 运行时 JIT 编译）；TilingData 29 字段
776+ 派生逻辑与 op_host/compressor_grad_tiling.cpp 一致。
777+ """
778+ p = _parse_case(case, device)
779+ B, H, D = p["B"], p["H"], p["D"]
780+ cr, coff = p["cr"], p["coff"]
781+ layoutId, ioDtype = p["layoutId"], p["ioDtype"]
782+ tokenSize, seqSize, outputRows = p["tokenSize"], p["seqSize"], p["outputRows"]
783+ sqs, sps = case.get("seqused_q"), case.get("start_pos")
784+ 
785+ if layoutId == 1:
786+ cuNpu = torch.tensor(case["seqlens_list_q"], device=device, dtype=torch.int32)
787+ else:
788+ cuNpu = None
789+ sqNpu = torch.tensor(sqs, device=device, dtype=torch.int32) if sqs else None
790+ spNpu = torch.tensor(sps, device=device, dtype=torch.int32) if sps else None
791+ 
792+ # ── TilingData（派生与 op_host/compressor_grad_tiling.cpp 一致）──
793+ cubeCoreNum, coreNum = 32, 64
794+ coffCoef = 2 // coff
795+ totalHeadDim = coff * D
796+ cmpSize = coff * cr * D
797+ cmpKvBatchStride = _ceil_div(seqSize, cr)
798+ xRows = tokenSize if layoutId == 1 else B * seqSize
799+ groupSize = D // 128
800+ groupNum = cubeCoreNum // groupSize
801+ dealScNum = 128 // cr
802+ groupDealScNum = dealScNum * coffCoef
803+ totalScNumPerRound = groupNum * groupDealScNum
804+ groupRowStride = groupDealScNum * cr + (coff - 1) * cr
805+ dbRowCnt = groupNum * groupRowStride
806+ tiling = CompressorGradTiling(
807+ batch_size=B, token_size=tokenSize, seq_size=seqSize,
808+ cmp_ratio=cr, hidden_size=H, head_dim=D,
809+ cube_core_num=cubeCoreNum, core_num=coreNum,
810+ total_head_dim=totalHeadDim, cmp_row_cnt=coff * cr, cmp_size=cmpSize,
811+ cmp_kv_batch_stride=cmpKvBatchStride, cmp_kv_rows=outputRows, x_rows=xRows,
812+ group_size=groupSize, group_num=groupNum,
813+ group_deal_sc_num=groupDealScNum, deal_sc_num=dealScNum,
814+ total_sc_num_per_round=totalScNumPerRound, db_row_cnt=dbRowCnt,
815+ group_row_stride=groupRowStride,
816+ coff_coef=coffCoef, cube_m_base_size=128 * coffCoef,
817+ d_deal_size=128 // coff, m_deal_size=128 * coff,
818+ dape_ws_size=groupNum * cmpSize * coffCoef,
819+ d_x_ws_size=2 * cubeCoreNum * 256 * H,
820+ d_w_weight_ws_size=groupNum * totalHeadDim * H,
821+ x_ws_size=2 * groupNum * groupRowStride * H,
822+ d_x_cache_ws_size=2 * cr * H,
823+ )
824+ 
825+ # ── 随机构造输入/中间量（不依赖正向输出）──
826+ torch.manual_seed(case["manual_seed"])
827+ xNpu = torch.randn(xRows, H, device=device, dtype=ioDtype)
828+ wkvInp = torch.randn(coff * D, H, device=device, dtype=ioDtype)
829+ wgateInp = torch.randn(coff * D, H, device=device, dtype=ioDtype)
830+ dcNpu = torch.randn(outputRows, D, device=device, dtype=ioDtype)
831+ kvNpu = torch.randn(outputRows, coff * cr, D, device=device, dtype=torch.float32)
832+ smNpu = torch.softmax(torch.randn(outputRows, coff * cr, D, device=device), dim=1)
833+ 
834+ # ── 输出与 workspace（精确大小，FP32 元素数）──
835+ # ⚠️ 全部用随机值初始化：kernel 未写区域将暴露为垃圾而非被 0 掩盖
836+ #（曾因 workspace=zeros 掩盖 d_ape 未写槽位读入问题）
837+ dX = torch.rand(xRows, H, device=device, dtype=ioDtype)
838+ dWkv = torch.rand(coff * D, H, device=device, dtype=ioDtype)
839+ dWgate = torch.rand(coff * D, H, device=device, dtype=ioDtype)
840+ dApe = torch.rand(cr * coff * D, device=device, dtype=torch.float32)
841+ wsSize = (tiling.dape_ws_size + tiling.d_x_ws_size + 2 * tiling.d_w_weight_ws_size +
842+ tiling.x_ws_size + tiling.d_x_cache_ws_size)
843+ workspace = torch.rand(wsSize, device=device, dtype=torch.float32)
844+ 
845+ # ── PyPTO 直跑（不经 aclnn/custom 包）──
846+ compressor_grad[None, cubeCoreNum,
847+ {"Coff": coff, "Layout": layoutId,
848+ "DataType": 0 if ioDtype == torch.bfloat16 else 1}](
849+ xNpu, wkvInp, wgateInp, dcNpu, smNpu, kvNpu,
850+ cuNpu, sqNpu, spNpu,
851+ dX, dWkv, dWgate, dApe, workspace, tiling)
852+ torch.npu.synchronize()
853+ 
854+ # ── golden（BSH 输入需 3 维；dc/sm/kv 需 (B, blocks, ...) 维）──
855+ xCpu = xNpu.cpu().view(B, seqSize, H) if layoutId == 0 else xNpu.cpu()
856+ smCpu = smNpu.cpu().view(B, -1, coff * cr, D) if layoutId == 0 else smNpu.cpu()
857+ kvCpu = kvNpu.cpu().view(B, -1, coff * cr, D) if layoutId == 0 else kvNpu.cpu()
858+ g_dx, g_dwkv, g_dwgate, g_dape = compressor_grad_golden(
859+ x=xCpu, wkv=wkvInp.cpu().float(), wgate=wgateInp.cpu().float(),
860+ d_cpm_kv=(dcNpu.cpu().view(B, -1, D) if layoutId == 0 else dcNpu.cpu()),
861+ softmax_score=smCpu, kv=kvCpu,
862+ cu_seqlens=cuNpu.cpu() if cuNpu is not None else None,
863+ seqused=sqNpu.cpu() if sqNpu is not None else None,
864+ start_pos=spNpu.cpu() if spNpu is not None else None,
865+ cmp_ratio=cr, coff=coff)
866+ 
867+ # ── 四输出比对 ──
868+ p_pct_thd = get_pct_thd(p["dataType"])
869+ checks = [
870+ ("d_wkv", g_dwkv.float(), dWkv.cpu().float(), p["dataType"], p_pct_thd),
871+ ("d_wgate", g_dwgate.float(), dWgate.cpu().float(), p["dataType"], p_pct_thd),
872+ ("d_ape", g_dape, dApe.cpu().float(), "float32", get_pct_thd("float32")),
873+ ("d_x", g_dx.float(), dX.cpu().float(), p["dataType"], p_pct_thd),
874+ ]
875+ statuses = []
876+ for name, exp, act, dt, thd in checks:
877+ r = check_one_output(name, exp, act, dt, True, 1, thd)
878+ statuses.append(f"{name}={r['status']}")
879+ status = "ERROR" if any("ERROR" in s for s in statuses) else \
880+ "FAIL" if any("FAIL" in s for s in statuses) else "PASS"
881+ return {"status": status, "detail": " ".join(statuses)}
882+ 
883+ 
884+def _pathway(request):
885+ """当前验证通路（--pathway 入参）。"""
886+ return request.config.getoption("--pathway")
887+ 
888+ 
889+@pytest.mark.parametrize(
890+ "case",
891+ [c for c in REDLINE_CASES if c.get("enable", True)],
892+ ids=[c["testcase_name"] for c in REDLINE_CASES if c.get("enable", True)],
893+)
894+def test_compressor_grad(case, request):
895+ """通路 1：autograd 正反向全链路 + 反向四输出校验
896+ + 正向输出校验（cmp_kv/softmax_score/kv vs compressor_forward 参考）。
897+ 
898+ 运行：pytest test_compressor_grad.py -k <case名子串> -q
899+ """
900+ if _pathway(request) != 1:
901+ pytest.skip(f"pathway={_pathway(request)}")
902+ device = f"npu:{DEVICE_ID}"
903+ torch.npu.set_device(device)
904+ result = run_case(case, device)
905+ assert result["status"] == "PASS", format_case_line(result)
906+ 
907+ 
908+@pytest.mark.parametrize(
909+ "case",
910+ [c for c in REDLINE_CASES if c.get("enable", True)],
911+ ids=[c["testcase_name"] for c in REDLINE_CASES if c.get("enable", True)],
912+)
913+def test_compressor_grad_backward(case, request):
914+ """通路 2：单反向（PyPTO 直跑，随机构造中间量，不经 custom 包）。
915+ 
916+ 运行：pytest test_compressor_grad.py --pathway 2 -k <case名子串> -q
917+ """
918+ if _pathway(request) != 2:
919+ pytest.skip(f"pathway={_pathway(request)}")
920+ result = run_backward_case(case, f"npu:{DEVICE_ID}")
921+ assert result["status"] == "PASS", result["detail"]
922+ 
923+ 
924+@pytest.mark.parametrize(
925+ "case",
926+ [c for c in REDLINE_CASES if c.get("enable", True)],
927+ ids=[c["testcase_name"] for c in REDLINE_CASES if c.get("enable", True)],
928+)
929+def test_compressor_grad_small_ops(case, request):
930+ """通路 3：小算子拼接 golden + pta 串联正反向（loss 只对有效压缩块）。
931+ 
932+ 运行：pytest test_compressor_grad.py --pathway 3 -k <case名子串> -q
933+ """
934+ if _pathway(request) != 3:
935+ pytest.skip(f"pathway={_pathway(request)}")
936+ result = run_small_ops_case(case, f"npu:{DEVICE_ID}")
937+ assert result["status"] == "PASS", result["detail"]
938+ 
939+ 
940+@pytest.mark.parametrize(
941+ "case",
942+ [c for c in REDLINE_CASES if c.get("enable", True)],
943+ ids=[c["testcase_name"] for c in REDLINE_CASES if c.get("enable", True)],
944+)
945+def test_compressor_grad_backward_direct(case, request):
946+ """通路 4：单反向——通过正向接口内部函数 _compressor_backward 直接调用
947+ （不经 autograd，经 custom 包/aclnn）；中间量来自真实正向（先直调一次
948+ _compressor_forward，与通路 1 的 golden 一致性原则相同）。
949+ 
950+ 运行：pytest test_compressor_grad.py --pathway 4 -k <case名子串> -q
951+ """
952+ if _pathway(request) != 4:
953+ pytest.skip(f"pathway={_pathway(request)}")
954+ device = f"npu:{DEVICE_ID}"
955+ torch.npu.set_device(device)
956+ p, totalValid, cuNpu, sps, sqs, seed = _prepare_case(case, device)
957+ b = _build_inputs(case, p, device, seed, totalValid, cuNpu, sqs, sps,
958+ state_init="rand")
959+ result = run_backward_direct_case(case, p, b, totalValid, device)
960+ assert result["status"] == "PASS", format_case_line(result)
961+ 
962+ 
963+@pytest.mark.parametrize(
964+ "case",
965+ [c for c in REDLINE_CASES if c.get("enable", True)],
966+ ids=[c["testcase_name"] for c in REDLINE_CASES if c.get("enable", True)],
967+)
968+def test_compressor_grad_forward_direct(case, request):
969+ """通路 5：单正向——通过正向接口内部函数 _compressor_forward 直接调用
970+ （不经 autograd，经 custom 包/aclnn），校验正向 3 输出
971+ （cmp_kv/softmax_score/kv）+ state_cache 4 项
972+ （kv_state_update/score_state_update/kv_state_origin/score_state_origin，
973+ 与正向算子 pytest 一致）。
974+ 
975+ 运行：pytest test_compressor_grad.py --pathway 5 -k <case名子串> -q
976+ """
977+ if _pathway(request) != 5:
978+ pytest.skip(f"pathway={_pathway(request)}")
979+ device = f"npu:{DEVICE_ID}"
980+ torch.npu.set_device(device)
981+ p, totalValid, cuNpu, sps, sqs, seed = _prepare_case(case, device)
982+ b = _build_inputs(case, p, device, seed, totalValid, cuNpu, sqs, sps,
983+ state_init="rand")
984+ result, _ = run_forward_direct_case(case, p, b, totalValid, device)
985+ assert result["status"] == "PASS", format_case_line(result)
Aattention/compressor_grad/tests/pytest/test_compressor_grad_cases.py+1314-0
Aattention/compressor_grad/tests/ut/CMakeLists.txt+16-0
@@ -0,0 +1,16 @@
1+# -----------------------------------------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# -----------------------------------------------------------------------------------------------------------
10+ 
11+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
12+foreach(SUB_DIR ${CURRENT_DIRS})
13+ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
14+ add_subdirectory(${SUB_DIR})
15+ endif()
16+endforeach()
Aattention/compressor_grad/tests/ut/op_host/CMakeLists.txt+13-0
@@ -0,0 +1,13 @@
1+# -----------------------------------------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# -----------------------------------------------------------------------------------------------------------
10+ 
11+if(UT_TEST_ALL OR OP_HOST_UT)
12+ add_modules_ut_sources(UT_NAME ${OP_TILING_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
13+endif()
Aattention/compressor_grad/tests/ut/op_host/test_compressor_grad_tiling.cpp+498-0
@@ -0,0 +1,498 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <gtest/gtest.h>
12+#include <iostream>
13+ 
14+#include "../../../op_host/compressor_grad_tiling.h"
15+#include "tiling_context_faker.h"
16+#include "tiling_case_executor.h"
17+ 
18+using namespace std;
19+ 
20+// Ascend950 SoC info
21+std::string CompressorGrad_tiling_A5SocInfo =
22+ "{\n"
23+ " \"hardware_info\": {\n"
24+ " \"BT_SIZE\": 0,\n"
25+ " \"load3d_constraints\": \"1\",\n"
26+ " \"Intrinsic_fix_pipe_l0c2out\": false,\n"
27+ " \"Intrinsic_data_move_l12ub\": true,\n"
28+ " \"Intrinsic_data_move_l0c2ub\": true,\n"
29+ " \"Intrinsic_data_move_out2l1_nd2nz\": false,\n"
30+ " \"4096\": 196608,\n"
31+ " \"L2_SIZE\": 201326592,\n"
32+ " \"L1_SIZE\": 524288,\n"
33+ " \"L0A_SIZE\": 65536,\n"
34+ " \"L0B_SIZE\": 65536,\n"
35+ " \"L0C_SIZE\": 131072,\n"
36+ " \"vector_core_cnt\": 40,\n"
37+ " \"cube_core_cnt\": 20,\n"
38+ " \"socVersion\": \"Ascend950\"\n"
39+ " }\n"
40+ "}";
41+ 
42+// ====================================================================
43+// Tiling key encoding (from compressor_grad.py CompressorGradTilingKey):
44+// key = coff * 16 + layout * 2 + dtype
45+// coff: 1 or 2 (2-bit field, values [1,2])
46+// layout: 0=BSH 1=TH (1-bit field)
47+// dtype: 0=BF16 1=FP16 (2-bit field)
48+//
49+// Verification against forward compressor tests:
50+// coff=2, BSH, BF16 → 2*16+0*2+0 = 32 ✓
51+// coff=2, BSH, FP16 → 2*16+0*2+1 = 34 ✓
52+// coff=1, BSH, BF16 → 1*16+0*2+0 = 16
53+// ====================================================================
54+ 
55+class CompressorGradTilingTest : public testing::Test {
56+protected:
57+ static void SetUpTestCase()
58+ {
59+ std::cout << "CompressorGradTilingTest SetUp" << std::endl;
60+ }
61+ static void TearDownTestCase()
62+ {
63+ std::cout << "CompressorGradTilingTest TearDown" << std::endl;
64+ }
65+};
66+ 
67+// ====================================================================
68+// BSH Layout Tiling Tests
69+// ====================================================================
70+ 
71+// C4A bf16: B=2, S=8, H=4096, D=512, coff=2, cmp_ratio=4
72+// tiling_key = 2*16 + 0*2 + 0 = 32
73+TEST_F(CompressorGradTilingTest, bsh_c4a_bf16)
74+{
75+ optiling::CompressorGradCompileInfo compileInfo = {};
76+ gert::TilingContextPara tilingContextPara(
77+ "CompressorGrad",
78+ {
79+ // Required inputs
80+ {{{2, 8, 4096}, {2, 8, 4096}}, ge::DT_BF16, ge::FORMAT_ND}, // x
81+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND}, // wkv
82+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND}, // wgate
83+ {{{2, 2, 512}, {2, 2, 512}}, ge::DT_BF16, ge::FORMAT_ND}, // d_cmp_kv
84+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score
85+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv
86+ // Optional inputs (empty for BSH)
87+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND}, // cu_seqlens
88+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND}, // seqused
89+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND}, // start_pos
90+ },
91+ {
92+ {{{2, 8, 4096}, {2, 8, 4096}}, ge::DT_BF16, ge::FORMAT_ND}, // d_x
93+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND}, // d_wkv
94+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND}, // d_wgate
95+ {{{4, 1024}, {4, 1024}}, ge::DT_FLOAT, ge::FORMAT_ND}, // d_ape
96+ },
97+ {
98+ {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
99+ {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
100+ },
101+ &compileInfo, "Ascend950", CompressorGrad_tiling_A5SocInfo, 4096);
102+ int64_t expectTilingKey = 1;
103+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey);
104+}
105+ 
106+// C4A fp16: B=2, S=8, H=4096, D=512, coff=2, cmp_ratio=4
107+// tiling_key = 2*16 + 0*2 + 1 = 34
108+TEST_F(CompressorGradTilingTest, bsh_c4a_fp16)
109+{
110+ optiling::CompressorGradCompileInfo compileInfo = {};
111+ gert::TilingContextPara tilingContextPara(
112+ "CompressorGrad",
113+ {
114+ {{{2, 8, 4096}, {2, 8, 4096}}, ge::DT_FLOAT16, ge::FORMAT_ND},
115+ {{{1024, 4096}, {1024, 4096}}, ge::DT_FLOAT16, ge::FORMAT_ND},
116+ {{{1024, 4096}, {1024, 4096}}, ge::DT_FLOAT16, ge::FORMAT_ND},
117+ {{{2, 2, 512}, {2, 2, 512}}, ge::DT_FLOAT16, ge::FORMAT_ND},
118+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
119+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
120+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
121+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
122+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
123+ },
124+ {
125+ {{{2, 8, 4096}, {2, 8, 4096}}, ge::DT_FLOAT16, ge::FORMAT_ND},
126+ {{{1024, 4096}, {1024, 4096}}, ge::DT_FLOAT16, ge::FORMAT_ND},
127+ {{{1024, 4096}, {1024, 4096}}, ge::DT_FLOAT16, ge::FORMAT_ND},
128+ {{{4, 1024}, {4, 1024}}, ge::DT_FLOAT, ge::FORMAT_ND},
129+ },
130+ {
131+ {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
132+ {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
133+ },
134+ &compileInfo, "Ascend950", CompressorGrad_tiling_A5SocInfo, 4096);
135+ int64_t expectTilingKey = 9;
136+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey);
137+}
138+ 
139+// C4Li bf16: B=2, S=8, H=2048, D=128, coff=2, cmp_ratio=4
140+// tiling_key = 2*16 + 0*2 + 0 = 32
141+TEST_F(CompressorGradTilingTest, bsh_c4li_bf16)
142+{
143+ optiling::CompressorGradCompileInfo compileInfo = {};
144+ gert::TilingContextPara tilingContextPara(
145+ "CompressorGrad",
146+ {
147+ {{{2, 8, 2048}, {2, 8, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
148+ {{{256, 2048}, {256, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
149+ {{{256, 2048}, {256, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
150+ {{{2, 2, 128}, {2, 2, 128}}, ge::DT_BF16, ge::FORMAT_ND},
151+ {{{2, 2, 8, 128}, {2, 2, 8, 128}}, ge::DT_FLOAT, ge::FORMAT_ND},
152+ {{{2, 2, 8, 128}, {2, 2, 8, 128}}, ge::DT_FLOAT, ge::FORMAT_ND},
153+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
154+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
155+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
156+ },
157+ {
158+ {{{2, 8, 2048}, {2, 8, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
159+ {{{256, 2048}, {256, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
160+ {{{256, 2048}, {256, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
161+ {{{4, 256}, {4, 256}}, ge::DT_FLOAT, ge::FORMAT_ND},
162+ },
163+ {
164+ {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
165+ {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
166+ },
167+ &compileInfo, "Ascend950", CompressorGrad_tiling_A5SocInfo, 4096);
168+ int64_t expectTilingKey = 1;
169+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey);
170+}
171+ 
172+// C128A bf16: B=2, S=16, H=4096, D=512, coff=1, cmp_ratio=128
173+// tiling_key = 1*16 + 0*2 + 0 = 16
174+TEST_F(CompressorGradTilingTest, bsh_c128a_bf16)
175+{
176+ optiling::CompressorGradCompileInfo compileInfo = {};
177+ gert::TilingContextPara tilingContextPara(
178+ "CompressorGrad",
179+ {
180+ {{{2, 16, 4096}, {2, 16, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
181+ {{{512, 4096}, {512, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
182+ {{{512, 4096}, {512, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
183+ {{{2, 1, 512}, {2, 1, 512}}, ge::DT_BF16, ge::FORMAT_ND},
184+ {{{2, 1, 128, 512}, {2, 1, 128, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
185+ {{{2, 1, 128, 512}, {2, 1, 128, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
186+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
187+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
188+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
189+ },
190+ {
191+ {{{2, 16, 4096}, {2, 16, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
192+ {{{512, 4096}, {512, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
193+ {{{512, 4096}, {512, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
194+ {{{128, 512}, {128, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
195+ },
196+ {
197+ {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(128)},
198+ {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},
199+ },
200+ &compileInfo, "Ascend950", CompressorGrad_tiling_A5SocInfo, 4096);
201+ int64_t expectTilingKey = 0;
202+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey);
203+}
204+ 
205+// Non-divisible S: B=1, S=6, H=2048, D=128, coff=2, cmp_ratio=4
206+// Sr = ceil(6/4) = 2, tiling_key = 2*16 + 0*2 + 0 = 32
207+TEST_F(CompressorGradTilingTest, bsh_non_divisible_s)
208+{
209+ optiling::CompressorGradCompileInfo compileInfo = {};
210+ gert::TilingContextPara tilingContextPara(
211+ "CompressorGrad",
212+ {
213+ {{{1, 6, 2048}, {1, 6, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
214+ {{{256, 2048}, {256, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
215+ {{{256, 2048}, {256, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
216+ {{{1, 2, 128}, {1, 2, 128}}, ge::DT_BF16, ge::FORMAT_ND},
217+ {{{1, 2, 8, 128}, {1, 2, 8, 128}}, ge::DT_FLOAT, ge::FORMAT_ND},
218+ {{{1, 2, 8, 128}, {1, 2, 8, 128}}, ge::DT_FLOAT, ge::FORMAT_ND},
219+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
220+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
221+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
222+ },
223+ {
224+ {{{1, 6, 2048}, {1, 6, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
225+ {{{256, 2048}, {256, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
226+ {{{256, 2048}, {256, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
227+ {{{4, 256}, {4, 256}}, ge::DT_FLOAT, ge::FORMAT_ND},
228+ },
229+ {
230+ {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
231+ {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
232+ },
233+ &compileInfo, "Ascend950", CompressorGrad_tiling_A5SocInfo, 4096);
234+ int64_t expectTilingKey = 1;
235+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey);
236+}
237+ 
238+// ====================================================================
239+// TH Layout Tiling Tests
240+// ====================================================================
241+ 
242+// TH C4A bf16: T=8, H=4096, B=2, cu_seqlens=[0,4,8], D=512, coff=2, cmp_ratio=4
243+// For backward: dimNum(dCmpKv)==dimNum(x)==2, softmaxScore/kv are 3D
244+// cmpKvRows = min(T, T/cmpRatio + B) = min(8, 2+2) = 4
245+// tiling_key = 2*16 + 1*2 + 0 = 35
246+TEST_F(CompressorGradTilingTest, th_c4a_bf16)
247+{
248+ optiling::CompressorGradCompileInfo compileInfo = {};
249+ gert::TilingContextPara tilingContextPara(
250+ "CompressorGrad",
251+ {
252+ {{{8, 4096}, {8, 4096}}, ge::DT_BF16, ge::FORMAT_ND}, // x (T, H)
253+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND}, // wkv
254+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND}, // wgate
255+ {{{4, 512}, {4, 512}}, ge::DT_BF16, ge::FORMAT_ND}, // d_cmp_kv (cmpKvRows, D)
256+ {{{2, 8, 512}, {2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmax_score (B, coff*cmpRatio, D)
257+ {{{2, 8, 512}, {2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND}, // kv (B, coff*cmpRatio, D)
258+ {{{3}, {3}}, ge::DT_INT32, ge::FORMAT_ND}, // cu_seqlens
259+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND}, // seqused
260+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND}, // start_pos
261+ },
262+ {
263+ {{{8, 4096}, {8, 4096}}, ge::DT_BF16, ge::FORMAT_ND}, // d_x (T, H)
264+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND}, // d_wkv
265+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND}, // d_wgate
266+ {{{4, 1024}, {4, 1024}}, ge::DT_FLOAT, ge::FORMAT_ND}, // d_ape
267+ },
268+ {
269+ {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
270+ {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
271+ },
272+ &compileInfo, "Ascend950", CompressorGrad_tiling_A5SocInfo, 4096);
273+ int64_t expectTilingKey = 5;
274+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey);
275+}
276+ 
277+// TH C4A fp16: T=8, H=4096, B=2, D=512, coff=2, cmp_ratio=4
278+// tiling_key = 2*16 + 1*2 + 1 = 36
279+TEST_F(CompressorGradTilingTest, th_c4a_fp16)
280+{
281+ optiling::CompressorGradCompileInfo compileInfo = {};
282+ gert::TilingContextPara tilingContextPara(
283+ "CompressorGrad",
284+ {
285+ {{{8, 4096}, {8, 4096}}, ge::DT_FLOAT16, ge::FORMAT_ND},
286+ {{{1024, 4096}, {1024, 4096}}, ge::DT_FLOAT16, ge::FORMAT_ND},
287+ {{{1024, 4096}, {1024, 4096}}, ge::DT_FLOAT16, ge::FORMAT_ND},
288+ {{{4, 512}, {4, 512}}, ge::DT_FLOAT16, ge::FORMAT_ND},
289+ {{{2, 8, 512}, {2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
290+ {{{2, 8, 512}, {2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
291+ {{{3}, {3}}, ge::DT_INT32, ge::FORMAT_ND},
292+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
293+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
294+ },
295+ {
296+ {{{8, 4096}, {8, 4096}}, ge::DT_FLOAT16, ge::FORMAT_ND},
297+ {{{1024, 4096}, {1024, 4096}}, ge::DT_FLOAT16, ge::FORMAT_ND},
298+ {{{1024, 4096}, {1024, 4096}}, ge::DT_FLOAT16, ge::FORMAT_ND},
299+ {{{4, 1024}, {4, 1024}}, ge::DT_FLOAT, ge::FORMAT_ND},
300+ },
301+ {
302+ {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
303+ {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
304+ },
305+ &compileInfo, "Ascend950", CompressorGrad_tiling_A5SocInfo, 4096);
306+ int64_t expectTilingKey = 13;
307+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey);
308+}
309+ 
310+// ====================================================================
311+// Error Cases — expected FAIL
312+// ====================================================================
313+ 
314+// Unsupported cmp_ratio=3 (interval check: [2, 128])
315+TEST_F(CompressorGradTilingTest, err_cmp_ratio_3)
316+{
317+ optiling::CompressorGradCompileInfo compileInfo = {};
318+ gert::TilingContextPara tilingContextPara(
319+ "CompressorGrad",
320+ {
321+ {{{2, 8, 4096}, {2, 8, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
322+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
323+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
324+ {{{2, 2, 512}, {2, 2, 512}}, ge::DT_BF16, ge::FORMAT_ND},
325+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
326+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
327+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
328+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
329+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
330+ },
331+ {
332+ {{{2, 8, 4096}, {2, 8, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
333+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
334+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
335+ {{{4, 1024}, {4, 1024}}, ge::DT_FLOAT, ge::FORMAT_ND},
336+ },
337+ {
338+ {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(3)},
339+ {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
340+ },
341+ &compileInfo, "Ascend950", CompressorGrad_tiling_A5SocInfo, 4096);
342+ ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED, 0);
343+}
344+ 
345+// Unsupported coff=3 (only 1 and 2 supported)
346+TEST_F(CompressorGradTilingTest, err_coff_3)
347+{
348+ optiling::CompressorGradCompileInfo compileInfo = {};
349+ gert::TilingContextPara tilingContextPara(
350+ "CompressorGrad",
351+ {
352+ {{{2, 8, 4096}, {2, 8, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
353+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
354+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
355+ {{{2, 2, 512}, {2, 2, 512}}, ge::DT_BF16, ge::FORMAT_ND},
356+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
357+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
358+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
359+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
360+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
361+ },
362+ {
363+ {{{2, 8, 4096}, {2, 8, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
364+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
365+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
366+ {{{4, 1024}, {4, 1024}}, ge::DT_FLOAT, ge::FORMAT_ND},
367+ },
368+ {
369+ {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
370+ {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(3)},
371+ },
372+ &compileInfo, "Ascend950", CompressorGrad_tiling_A5SocInfo, 4096);
373+ ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED, 0);
374+}
375+ 
376+// Unsupported headDim=256 (only 128 and 512 supported)
377+TEST_F(CompressorGradTilingTest, err_head_dim_256)
378+{
379+ optiling::CompressorGradCompileInfo compileInfo = {};
380+ gert::TilingContextPara tilingContextPara(
381+ "CompressorGrad",
382+ {
383+ {{{2, 8, 2048}, {2, 8, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
384+ {{{512, 2048}, {512, 2048}}, ge::DT_BF16, ge::FORMAT_ND}, // coff*headDim = 2*256 = 512
385+ {{{512, 2048}, {512, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
386+ {{{2, 2, 256}, {2, 2, 256}}, ge::DT_BF16, ge::FORMAT_ND},
387+ {{{2, 2, 8, 256}, {2, 2, 8, 256}}, ge::DT_FLOAT, ge::FORMAT_ND},
388+ {{{2, 2, 8, 256}, {2, 2, 8, 256}}, ge::DT_FLOAT, ge::FORMAT_ND},
389+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
390+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
391+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
392+ },
393+ {
394+ {{{2, 8, 2048}, {2, 8, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
395+ {{{512, 2048}, {512, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
396+ {{{512, 2048}, {512, 2048}}, ge::DT_BF16, ge::FORMAT_ND},
397+ {{{4, 512}, {4, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
398+ },
399+ {
400+ {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
401+ {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
402+ },
403+ &compileInfo, "Ascend950", CompressorGrad_tiling_A5SocInfo, 4096);
404+ ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED, 0);
405+}
406+ 
407+// Unsupported hiddenSize=3000 (not 512-aligned, must be in [1024, 10240])
408+TEST_F(CompressorGradTilingTest, err_hidden_size_not_aligned)
409+{
410+ optiling::CompressorGradCompileInfo compileInfo = {};
411+ gert::TilingContextPara tilingContextPara(
412+ "CompressorGrad",
413+ {
414+ {{{2, 8, 3000}, {2, 8, 3000}}, ge::DT_BF16, ge::FORMAT_ND},
415+ {{{1024, 3000}, {1024, 3000}}, ge::DT_BF16, ge::FORMAT_ND},
416+ {{{1024, 3000}, {1024, 3000}}, ge::DT_BF16, ge::FORMAT_ND},
417+ {{{2, 2, 512}, {2, 2, 512}}, ge::DT_BF16, ge::FORMAT_ND},
418+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
419+ {{{2, 2, 8, 512}, {2, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
420+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
421+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
422+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
423+ },
424+ {
425+ {{{2, 8, 3000}, {2, 8, 3000}}, ge::DT_BF16, ge::FORMAT_ND},
426+ {{{1024, 3000}, {1024, 3000}}, ge::DT_BF16, ge::FORMAT_ND},
427+ {{{1024, 3000}, {1024, 3000}}, ge::DT_BF16, ge::FORMAT_ND},
428+ {{{4, 1024}, {4, 1024}}, ge::DT_FLOAT, ge::FORMAT_ND},
429+ },
430+ {
431+ {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
432+ {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
433+ },
434+ &compileInfo, "Ascend950", CompressorGrad_tiling_A5SocInfo, 4096);
435+ ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED, 0);
436+}
437+ 
438+// Unsupported hiddenSize=512 (below MIN_HIDDEN_SIZE=1024)
439+TEST_F(CompressorGradTilingTest, err_hidden_size_too_small)
440+{
441+ optiling::CompressorGradCompileInfo compileInfo = {};
442+ gert::TilingContextPara tilingContextPara(
443+ "CompressorGrad",
444+ {
445+ {{{2, 8, 512}, {2, 8, 512}}, ge::DT_BF16, ge::FORMAT_ND},
446+ {{{256, 512}, {256, 512}}, ge::DT_BF16, ge::FORMAT_ND},
447+ {{{256, 512}, {256, 512}}, ge::DT_BF16, ge::FORMAT_ND},
448+ {{{2, 2, 128}, {2, 2, 128}}, ge::DT_BF16, ge::FORMAT_ND},
449+ {{{2, 2, 8, 128}, {2, 2, 8, 128}}, ge::DT_FLOAT, ge::FORMAT_ND},
450+ {{{2, 2, 8, 128}, {2, 2, 8, 128}}, ge::DT_FLOAT, ge::FORMAT_ND},
451+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
452+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
453+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
454+ },
455+ {
456+ {{{2, 8, 512}, {2, 8, 512}}, ge::DT_BF16, ge::FORMAT_ND},
457+ {{{256, 512}, {256, 512}}, ge::DT_BF16, ge::FORMAT_ND},
458+ {{{256, 512}, {256, 512}}, ge::DT_BF16, ge::FORMAT_ND},
459+ {{{4, 256}, {4, 256}}, ge::DT_FLOAT, ge::FORMAT_ND},
460+ },
461+ {
462+ {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
463+ {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
464+ },
465+ &compileInfo, "Ascend950", CompressorGrad_tiling_A5SocInfo, 4096);
466+ ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED, 0);
467+}
468+ 
469+// Empty tensor: x has zero-size shape → should fail (backward does not support empty tensors)
470+TEST_F(CompressorGradTilingTest, err_empty_x_tensor)
471+{
472+ optiling::CompressorGradCompileInfo compileInfo = {};
473+ gert::TilingContextPara tilingContextPara(
474+ "CompressorGrad",
475+ {
476+ {{{0, 8, 4096}, {0, 8, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
477+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
478+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
479+ {{{0, 2, 512}, {0, 2, 512}}, ge::DT_BF16, ge::FORMAT_ND},
480+ {{{0, 2, 8, 512}, {0, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
481+ {{{0, 2, 8, 512}, {0, 2, 8, 512}}, ge::DT_FLOAT, ge::FORMAT_ND},
482+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
483+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
484+ {{{}, {}}, ge::DT_INT32, ge::FORMAT_ND},
485+ },
486+ {
487+ {{{0, 8, 4096}, {0, 8, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
488+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
489+ {{{1024, 4096}, {1024, 4096}}, ge::DT_BF16, ge::FORMAT_ND},
490+ {{{4, 1024}, {4, 1024}}, ge::DT_FLOAT, ge::FORMAT_ND},
491+ },
492+ {
493+ {"cmp_ratio", Ops::Transformer::AnyValue::CreateFrom<int64_t>(4)},
494+ {"coff", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2)},
495+ },
496+ &compileInfo, "Ascend950", CompressorGrad_tiling_A5SocInfo, 4096);
497+ ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED, 0);
498+}
Mtorch_extension/cann_ops_transformer/ops/compressor.py+175-56
@@ -7,19 +7,13 @@
7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8# See LICENSE in the root of the software repository for the full text of the License.8# See LICENSE in the root of the software repository for the full text of the License.
9# -----------------------------------------------------------------------------------------------------------9# -----------------------------------------------------------------------------------------------------------
10-from enum import IntEnum10+from typing import Optional, Tuple
11-from typing import Optional
12import torch11import torch
13from torch.library import impl12from torch.library import impl
14from cann_ops_transformer.op_builder.builder import OpBuilder13from cann_ops_transformer.op_builder.builder import OpBuilder
15from cann_ops_transformer.op_builder.builder import AS_LIBRARY14from cann_ops_transformer.op_builder.builder import AS_LIBRARY
16 15 
17 16 
18-class CacheMode(IntEnum):
19- LINEAR_BUFFER = 1
20- RING_BUFFER = 2
21- 
22- 
23class CompressorOpBuilder(OpBuilder):17class CompressorOpBuilder(OpBuilder):
24 def __init__(self):18 def __init__(self):
25 super(CompressorOpBuilder, self).__init__("compressor")19 super(CompressorOpBuilder, self).__init__("compressor")
@@ -30,75 +24,39 @@ class CompressorOpBuilder(OpBuilder):
30 24 
31 def schema(self):25 def schema(self):
32 """PyTorch operator signature."""26 """PyTorch operator signature."""
33- return (27+ pass
34- "compressor(Tensor x, Tensor wkv, Tensor wgate, Tensor(a!) state_cache, "
35- "Tensor ape, "
36- "int cmp_ratio=4, *, "
37- "Tensor? state_block_table=None, Tensor? cu_seqlens=None, "
38- "Tensor? seqused=None, Tensor? start_pos=None, "
39- "int coff=1, int cache_mode=1) -> Tensor"
40- )
41 28 
42 def register_meta(self):29 def register_meta(self):
43 """30 """
44 Registers the Meta implementation (Shape/Dtype inference).31 Registers the Meta implementation (Shape/Dtype inference).
45 Essential for Autograd and FakeTensor support.32 Essential for Autograd and FakeTensor support.
46 """33 """
47- 34+ pass
48- @impl(AS_LIBRARY, self.name, "Meta")
49- def compressor_meta(
50- x,
51- wkv,
52- wgate,
53- state_cache,
54- ape,
55- cmp_ratio=4,
56- *,
57- state_block_table=None,
58- cu_seqlens=None,
59- seqused=None,
60- start_pos=None,
61- coff=1,
62- cache_mode=CacheMode.LINEAR_BUFFER,
63- ):
64- d = wkv.size(0) // coff
65- if x.dim() == 3:
66- b = x.size(0)
67- s = x.size(1)
68- sr = (s + cmp_ratio - 1) // cmp_ratio
69- cmp_kv_size = (b, sr, d)
70- else:
71- t = x.size(0)
72- sr = (t + cmp_ratio - 1) // cmp_ratio
73- cmp_kv_size = (sr, d)
74- 
75- return torch.empty(cmp_kv_size, dtype=x.dtype, device="meta")
76 35 
77 36 
78compressor_op_builder = CompressorOpBuilder()37compressor_op_builder = CompressorOpBuilder()
38+op_module = compressor_op_builder.load()
79 39 
80 40 
81-@impl(AS_LIBRARY, compressor_op_builder.name, "PrivateUse1")41+# ===========================================================================
82-def compressor(42+# Register compressor forward
43+# ===========================================================================
44+@torch.library.custom_op("cann_ops_transformer::_compressor_forward", mutates_args=(), device_types="npu")
45+def _compressor_forward(
83 x: torch.Tensor,46 x: torch.Tensor,
84 wkv: torch.Tensor,47 wkv: torch.Tensor,
85 wgate: torch.Tensor,48 wgate: torch.Tensor,
86 state_cache: torch.Tensor,49 state_cache: torch.Tensor,
87 ape: torch.Tensor,50 ape: torch.Tensor,
88- cmp_ratio: int = 4,
89- *,
90 state_block_table: Optional[torch.Tensor] = None,51 state_block_table: Optional[torch.Tensor] = None,
91 cu_seqlens: Optional[torch.Tensor] = None,52 cu_seqlens: Optional[torch.Tensor] = None,
92 seqused: Optional[torch.Tensor] = None,53 seqused: Optional[torch.Tensor] = None,
93 start_pos: Optional[torch.Tensor] = None,54 start_pos: Optional[torch.Tensor] = None,
55+ cmp_ratio: int = 4,
94 coff: Optional[int] = 1,56 coff: Optional[int] = 1,
95- cache_mode: Optional[CacheMode] = CacheMode.LINEAR_BUFFER,57+ cache_mode: Optional[int] = 1,
96-) -> torch.tensor:58+ grad_enabled: Optional[bool] = False
97- """59+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
98- dispatcher implementation for NPU.
99- 'PrivateUse1' is the combine key for custom NPU backends.
100- """
101- op_module = compressor_op_builder.load()
102 return op_module.compressor(60 return op_module.compressor(
103 x,61 x,
104 wkv,62 wkv,
@@ -112,7 +70,168 @@ def compressor(
112 start_pos,70 start_pos,
113 coff,71 coff,
114 cache_mode,72 cache_mode,
73+ grad_enabled
115 )74 )
116 75 
117 76 
118-compressor.CacheMode = CacheMode77+@torch.library.register_fake("cann_ops_transformer::_compressor_forward")
78+def _compressor_forward_fake(
79+ x: torch.Tensor,
80+ wkv: torch.Tensor,
81+ wgate: torch.Tensor,
82+ state_cache: torch.Tensor,
83+ ape: torch.Tensor,
84+ state_block_table: Optional[torch.Tensor] = None,
85+ cu_seqlens: Optional[torch.Tensor] = None,
86+ seqused: Optional[torch.Tensor] = None,
87+ start_pos: Optional[torch.Tensor] = None,
88+ cmp_ratio: int = 4,
89+ coff: Optional[int] = 1,
90+ cache_mode: Optional[int] = 1,
91+ grad_enabled: Optional[bool] = False
92+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
93+ d = wkv.size(0) // coff
94+ coff_cmp = coff * cmp_ratio
95+ if x.dim() == 3:
96+ b = x.size(0)
97+ s = x.size(1)
98+ cmp_size = (s + cmp_ratio - 1) // cmp_ratio
99+ cmp_kv_size = (b, cmp_size, d)
100+ softmax_score_size = (b, cmp_size, coff_cmp, d)
101+ kv_size = (b, cmp_size, coff_cmp, d)
102+ else:
103+ b_size = cu_seqlens.size(0) - 1
104+ t = x.size(0)
105+ cmp_size = min(t, t // cmp_ratio + b_size)
106+ cmp_kv_size = (cmp_size, d)
107+ softmax_score_size = (cmp_size, coff_cmp, d)
108+ kv_size = (cmp_size, coff_cmp, d)
109+
110+ cmp_kv_out = torch.empty(cmp_kv_size, dtype=x.dtype, device=x.device)
111+ softmax_score_out = torch.empty(
112+ softmax_score_size, dtype=torch.float32, device=x.device
113+ )
114+ kv_out = torch.empty(kv_size, dtype=torch.float32, device=x.device)
115+ return (cmp_kv_out, softmax_score_out, kv_out)
116+ 
117+ 
118+# ===========================================================================
119+# Register compressor backward
120+# ===========================================================================
121+@torch.library.custom_op("cann_ops_transformer::_compressor_backward", mutates_args=(), device_types="npu")
122+def _compressor_backward(
123+ d_cmp_kv: torch.Tensor,
124+ x: torch.Tensor,
125+ wkv: torch.Tensor,
126+ wgate: torch.Tensor,
127+ softmax_score: torch.Tensor,
128+ kv: torch.Tensor,
129+ cu_seqlens: Optional[torch.Tensor] = None,
130+ seqused: Optional[torch.Tensor] = None,
131+ start_pos: Optional[torch.Tensor] = None,
132+ cmp_ratio: int = 4,
133+ coff: Optional[int] = 1,
134+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
135+ d_x = torch.empty_like(x)
136+ d_wkv = torch.empty_like(wkv)
137+ d_wgate = torch.empty_like(wgate)
138+ d_ape_size = (cmp_ratio, wkv.size(0))
139+ d_ape = torch.empty(d_ape_size, dtype=torch.float32, device=x.device)
140+ return op_module.compressor_backward(
141+ d_cmp_kv, x, wkv, wgate, softmax_score, kv,
142+ cu_seqlens, seqused, start_pos, cmp_ratio, coff
143+ )
144+ 
145+ 
146+@torch.library.register_fake("cann_ops_transformer::_compressor_backward")
147+def _compressor_backward_fake(
148+ d_cmp_kv: torch.Tensor,
149+ x: torch.Tensor,
150+ wkv: torch.Tensor,
151+ wgate: torch.Tensor,
152+ softmax_score: torch.Tensor,
153+ kv: torch.Tensor,
154+ cu_seqlens: Optional[torch.Tensor] = None,
155+ seqused: Optional[torch.Tensor] = None,
156+ start_pos: Optional[torch.Tensor] = None,
157+ cmp_ratio: int = 4,
158+ coff: Optional[int] = 1,
159+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
160+ d_x = torch.empty_like(x)
161+ d_wkv = torch.empty_like(wkv)
162+ d_wgate = torch.empty_like(wgate)
163+ d_ape_size = (cmp_ratio, wkv.size(0))
164+ d_ape = torch.empty(d_ape_size, dtype=torch.float32, device=x.device)
165+ return (d_x, d_wkv, d_wgate, d_ape)
166+ 
167+ 
168+# ===========================================================================
169+# Register AutoGrad
170+# ===========================================================================
171+def setup_context(ctx, inputs, output):
172+ x, wkv, wgate = inputs[:3]
173+ cu_seqlens, seqused, start_pos, cmp_ratio, coff = inputs[6:11]
174+ 
175+ cmp_kv, softmax_score, kv = output
176+ ctx.save_for_backward(
177+ x, wkv, wgate, cu_seqlens, seqused, start_pos, softmax_score, kv
178+ )
179+ ctx.cmp_ratio = cmp_ratio
180+ ctx.coff = coff
181+ 
182+ 
183+def backward(ctx, dout, *grads):
184+ """
185+ Args:
186+ dout: d_cmp_kv
187+ """
188+ x, wkv, wgate, cu_seqlens, seqused, start_pos, softmax_score, kv = ctx.saved_tensors
189+ cmp_ratio = ctx.cmp_ratio
190+ coff = ctx.coff
191+ 
192+ d_x, d_wkv, d_wgate, d_ape = _compressor_backward(
193+ dout, x, wkv, wgate, softmax_score, kv,
194+ cu_seqlens, seqused, start_pos, cmp_ratio, coff
195+ )
196+ return d_x, d_wkv, d_wgate, None, d_ape, *((None, ) * 8)
197+
198+ 
199+_compressor_forward.register_autograd(backward, setup_context=setup_context)
200+ 
201+ 
202+def compressor(
203+ x: torch.Tensor,
204+ wkv: torch.Tensor,
205+ wgate: torch.Tensor,
206+ state_cache: torch.Tensor,
207+ ape: torch.Tensor,
208+ cmp_ratio: int = 4,
209+ *,
210+ state_block_table: Optional[torch.Tensor] = None,
211+ cu_seqlens: Optional[torch.Tensor] = None,
212+ seqused: Optional[torch.Tensor] = None,
213+ start_pos: Optional[torch.Tensor] = None,
214+ coff: Optional[int] = 1,
215+ cache_mode: Optional[int] = 1,
216+) -> Tuple[torch.Tensor]:
217+ """
218+ dispatcher implementation for NPU.
219+ 'PrivateUse1' is the combine key for custom NPU backends.
220+ """
221+ grad_enabled = x.requires_grad
222+ cmp_kv, softmax_score, kv = _compressor_forward(
223+ x,
224+ wkv,
225+ wgate,
226+ state_cache,
227+ ape,
228+ state_block_table,
229+ cu_seqlens,
230+ seqused,
231+ start_pos,
232+ cmp_ratio,
233+ coff,
234+ cache_mode,
235+ grad_enabled
236+ )
237+ return cmp_kv
Mtorch_extension/cann_ops_transformer/ops/csrc/compressor.cpp+70-17
@@ -46,12 +46,17 @@ std::vector<bool> IsContiguousAxes(const at::Tensor &tensor)
46 return result;46 return result;
47}47}
48 48 
49-at::Tensor ConstructCompressorOutputTensor(const at::Tensor &x, const at::Tensor &wkv,49+std::tuple<at::Tensor, at::Tensor, at::Tensor> ConstructCompressorOutputTensor(
50- const c10::optional<at::Tensor> &cuSeqlens, int64_t cmpRatio, int64_t coff)50+ const at::Tensor &x, const at::Tensor &wkv, const at::Tensor &ape,
51+ const c10::optional<at::Tensor> &cuSeqlens, int64_t cmpRatio, int64_t coff)
51{52{
52 auto xDim = x.dim();53 auto xDim = x.dim();
53 at::SmallVector<int64_t, MAX_DIM_SIZE> cmpKvSize;54 at::SmallVector<int64_t, MAX_DIM_SIZE> cmpKvSize;
55+ at::SmallVector<int64_t, MAX_DIM_SIZE> softmaxScoreSize;
56+ at::SmallVector<int64_t, MAX_DIM_SIZE> kvSize;
54 at::Tensor cmpKv;57 at::Tensor cmpKv;
58+ at::Tensor softmaxScore;
59+ at::Tensor kv;
55 int64_t cmpS = 0;60 int64_t cmpS = 0;
56 int64_t bSize = 0;61 int64_t bSize = 0;
57 62 
@@ -64,30 +69,37 @@ at::Tensor ConstructCompressorOutputTensor(const at::Tensor &x, const at::Tensor
64 if (xDim == DIM_THREE) {69 if (xDim == DIM_THREE) {
65 cmpS = (x.size(1) + cmpRatio - 1) / cmpRatio;70 cmpS = (x.size(1) + cmpRatio - 1) / cmpRatio;
66 cmpKvSize = {x.size(0), cmpS, wkv.size(0) / coff};71 cmpKvSize = {x.size(0), cmpS, wkv.size(0) / coff};
72+ softmaxScoreSize = {x.size(0), cmpS, coff * cmpRatio, wkv.size(0) / coff};
73+ kvSize = {x.size(0), cmpS, coff * cmpRatio, wkv.size(0) / coff};
67 } else {74 } else {
68 TORCH_CHECK(cuSeqlens.has_value(), "Check cu_seqlens != nullptr failed");75 TORCH_CHECK(cuSeqlens.has_value(), "Check cu_seqlens != nullptr failed");
69 bSize = cuSeqlens->size(0) - 1;76 bSize = cuSeqlens->size(0) - 1;
70 cmpS = std::min(x.size(0), x.size(0) / cmpRatio + bSize);77 cmpS = std::min(x.size(0), x.size(0) / cmpRatio + bSize);
71 cmpKvSize = {cmpS, wkv.size(0) / coff};78 cmpKvSize = {cmpS, wkv.size(0) / coff};
79+ softmaxScoreSize = {cmpS, coff * cmpRatio, wkv.size(0) / coff};
80+ kvSize = {cmpS, coff * cmpRatio, wkv.size(0) / coff};
72 }81 }
73 82 
74 cmpKv = at::empty(cmpKvSize, x.options().dtype(x.dtype()));83 cmpKv = at::empty(cmpKvSize, x.options().dtype(x.dtype()));
75- return cmpKv;84+ softmaxScore = at::empty(softmaxScoreSize, ape.options().dtype(ape.dtype()));
85+ kv = at::empty(kvSize, ape.options().dtype(ape.dtype()));
86+ return std::tuple<at::Tensor, at::Tensor, at::Tensor>(cmpKv, softmaxScore, kv);
76}87}
77 88 
78-at::Tensor Compressor(const at::Tensor &x, const at::Tensor &wkv, const at::Tensor &wgate, at::Tensor &stateCache,89+std::tuple<at::Tensor, at::Tensor, at::Tensor> Compressor(
79- const at::Tensor &ape, int64_t cmpRatio, const c10::optional<at::Tensor> &stateBlockTable,90+ const at::Tensor &x, const at::Tensor &wkv, const at::Tensor &wgate,
80- const c10::optional<at::Tensor> &cuSeqlens, const c10::optional<at::Tensor> &seqused,91+ at::Tensor &stateCache, const at::Tensor &ape, int64_t cmpRatio,
81- const c10::optional<at::Tensor> &startPos, int64_t coff, int64_t cacheMode)92+ const c10::optional<at::Tensor> &stateBlockTable,
93+ const c10::optional<at::Tensor> &cuSeqlens,
94+ const c10::optional<at::Tensor> &seqused, const c10::optional<at::Tensor> &startPos,
95+ int64_t coff, int64_t cacheMode, bool gradEnabled)
82{96{
83 TORCH_CHECK(x.defined(), "Check x != nullptr failed");97 TORCH_CHECK(x.defined(), "Check x != nullptr failed");
84 auto xDim = x.dim();98 auto xDim = x.dim();
85 TORCH_CHECK(xDim == DIM_TWO || xDim == DIM_THREE, "x dim num[", xDim, "] should be 2 or 3");99 TORCH_CHECK(xDim == DIM_TWO || xDim == DIM_THREE, "x dim num[", xDim, "] should be 2 or 3");
86 100 
87 TORCH_CHECK(cmpRatio > VALUE_0, "cmp_ratio should be greater than 0");101 TORCH_CHECK(cmpRatio > VALUE_0, "cmp_ratio should be greater than 0");
88- 102+ auto [cmpKv, softmaxScore, kv] = ConstructCompressorOutputTensor(x, wkv, ape, cuSeqlens, cmpRatio, coff);
89- at::Tensor cmpKv = ConstructCompressorOutputTensor(x, wkv, cuSeqlens, cmpRatio, coff);
90- 
91 auto stateCacheDim = stateCache.dim();103 auto stateCacheDim = stateCache.dim();
92 TORCH_CHECK(stateCacheDim == DIM_THREE, "state_cache dim num[", stateCacheDim, "] should be 3");104 TORCH_CHECK(stateCacheDim == DIM_THREE, "state_cache dim num[", stateCacheDim, "] should be 3");
93 105 
@@ -95,15 +107,54 @@ at::Tensor Compressor(const at::Tensor &x, const at::Tensor &wkv, const at::Tens
95 int64_t stateCacheStrideDim0 = stateCache.stride(0);107 int64_t stateCacheStrideDim0 = stateCache.stride(0);
96 108 
97 ACLNN_CMD(aclnnCompressor, x, wkv, wgate, stateCache, ape, stateBlockTable, cuSeqlens, seqused, startPos, cmpRatio,109 ACLNN_CMD(aclnnCompressor, x, wkv, wgate, stateCache, ape, stateBlockTable, cuSeqlens, seqused, startPos, cmpRatio,
98- coff, cacheMode, stateCacheStrideDim0, cmpKv);110+ coff, cacheMode, stateCacheStrideDim0, gradEnabled, cmpKv, softmaxScore, kv);
99 111 
100- return cmpKv;112+ return std::tuple<at::Tensor, at::Tensor, at::Tensor>(cmpKv, softmaxScore, kv);
101}113}
102 114 
103-at::Tensor CompressorMeta(const at::Tensor &x, const at::Tensor &wkv, const at::Tensor &wgate, at::Tensor &stateCache,115+std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> CompressorBackward(
104- const at::Tensor &ape, int64_t cmpRatio, const c10::optional<at::Tensor> &stateBlockTable,116+ at::Tensor &dCmpKv, const at::Tensor &x, const at::Tensor &wkv,
105- const c10::optional<at::Tensor> &cuSeqlens, const c10::optional<at::Tensor> &seqused,117+ const at::Tensor &wgate, const at::Tensor &softmaxScore, const at::Tensor &kv,
106- const c10::optional<at::Tensor> &startPos, int64_t coff, int64_t cacheMode)118+ const c10::optional<at::Tensor> &cuSeqlens, const c10::optional<at::Tensor> &seqused,
119+ const c10::optional<at::Tensor> &startPos, int64_t cmpRatio, int64_t coff)
120+{
121+ auto xDim = x.dim();
122+ at::SmallVector<int64_t, MAX_DIM_SIZE> dxSize;
123+ at::SmallVector<int64_t, MAX_DIM_SIZE> dWkvSize;
124+ at::SmallVector<int64_t, MAX_DIM_SIZE> dWgateSize;
125+ at::SmallVector<int64_t, MAX_DIM_SIZE> dApeSize;
126+ at::Tensor dx;
127+ at::Tensor dWkv;
128+ at::Tensor dWgate;
129+ at::Tensor dApe;
130+ 
131+ TORCH_CHECK(x.defined(), "Check x != nullptr failed");
132+ TORCH_CHECK(xDim == DIM_TWO || xDim == DIM_THREE, "x dim num[", xDim, "] should be 2 or 3");
133+ 
134+ TORCH_CHECK(cmpRatio > VALUE_0, "cmp_ratio should be greater than 0");
135+
136+ dxSize = x.sizes().vec();
137+ dx = at::empty(dxSize, x.options().dtype(x.dtype()));
138+ dWkvSize = wkv.sizes().vec();
139+ dWkv = at::empty(dWkvSize, wkv.options().dtype(wkv.dtype()));
140+ dWgateSize = wgate.sizes().vec();
141+ dWgate = at::empty(dWgateSize, wgate.options().dtype(wgate.dtype()));
142+ dApeSize = {cmpRatio, wkv.size(0)};
143+ dApe = at::empty(dApeSize, kv.options().dtype(kv.dtype()));
144+ 
145+ ACLNN_CMD(aclnnCompressorGrad, x, wkv, wgate, dCmpKv, softmaxScore, kv, cuSeqlens, seqused,
146+ startPos, cmpRatio, coff, dx, dWkv, dWgate, dApe);
147+ 
148+ return std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor>(dx, dWkv, dWgate, dApe);
149+}
150+ 
151+std::tuple<at::Tensor, at::Tensor, at::Tensor> CompressorMeta(
152+ const at::Tensor &x, const at::Tensor &wkv, const at::Tensor &wgate,
153+ at::Tensor &stateCache, const at::Tensor &ape, int64_t cmpRatio,
154+ const c10::optional<at::Tensor> &stateBlockTable,
155+ const c10::optional<at::Tensor> &cuSeqlens,
156+ const c10::optional<at::Tensor> &seqused, const c10::optional<at::Tensor> &startPos,
157+ int64_t coff, int64_t cacheMode)
107{158{
108 TORCH_CHECK(x.defined(), "Check x != nullptr failed");159 TORCH_CHECK(x.defined(), "Check x != nullptr failed");
109 auto xDim = x.dim();160 auto xDim = x.dim();
@@ -111,11 +162,13 @@ at::Tensor CompressorMeta(const at::Tensor &x, const at::Tensor &wkv, const at::
111 162 
112 TORCH_CHECK(cmpRatio > VALUE_0, "cmp_ratio should be greater than 0");163 TORCH_CHECK(cmpRatio > VALUE_0, "cmp_ratio should be greater than 0");
113 164 
114- return ConstructCompressorOutputTensor(x, wkv, cuSeqlens, cmpRatio, coff);165+ auto [cmpKv, softmaxScore, kv] = ConstructCompressorOutputTensor(x, wkv, ape, cuSeqlens, cmpRatio, coff);
166+ return std::tuple<at::Tensor, at::Tensor, at::Tensor>(cmpKv, softmaxScore, kv);
115}167}
116 168 
117PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)169PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
118{170{
119 m.def("compressor", &Compressor, "compressor");171 m.def("compressor", &Compressor, "compressor");
172+ m.def("compressor_backward", &CompressorBackward, "compressor_backward");
120}173}
121} // namespace op_api174} // namespace op_api
Mtorch_extension/cann_ops_transformer/ops/graph_convert/graph_convert_compressor.py+3-1
@@ -61,12 +61,13 @@ if _TORCHAIR_AVAILABLE:
61 *,61 *,
62 coff: int = 1,62 coff: int = 1,
63 cache_mode: int = 1,63 cache_mode: int = 1,
64+ grad_enabled: bool = False,
64 state_cache_stride_dim0: int = 0,65 state_cache_stride_dim0: int = 0,
65 ):66 ):
66 result = x.new_empty(x.size())67 result = x.new_empty(x.size())
67 return result68 return result
68 69 
69- @register_fx_node_ge_converter(torch.ops.cann_ops_transformer.compressor.default)70+ @register_fx_node_ge_converter(torch.ops.cann_ops_transformer._compressor_forward.default)
70 def convert_compressor(71 def convert_compressor(
71 x: Tensor,72 x: Tensor,
72 wkv: Tensor,73 wkv: Tensor,
@@ -81,6 +82,7 @@ if _TORCHAIR_AVAILABLE:
81 *,82 *,
82 coff: int = 1,83 coff: int = 1,
83 cache_mode: int = 1,84 cache_mode: int = 1,
85+ grad_enabled: bool = False,
84 state_cache_stride_dim0: int = 0,86 state_cache_stride_dim0: int = 0,
85 ):87 ):
86 raise AssertionError(88 raise AssertionError(