update dsv4 metadata param check-ops-transformer-AtomGit

Mattention/lightning_indexer_v2_metadata/op_host/lightning_indexer_v2_metadata_check.h+21-5

	@@ -91,9 +91,10 @@ int64_t GetKeyBatchSizeLiV2(int64_t batchSize, const aclTensor *cuSeqlensKOption
91	return batchSize;	91	return batchSize;
92	}	92	}
93		93
94	-aclnnStatus CheckSingleParamLiV2(int64_t numHeadsQ, int64_t numHeadsK, int64_t topk, int64_t batchSize,	94	+aclnnStatus CheckSingleParamLiV2(int64_t numHeadsQ, int64_t numHeadsK, int64_t headDim, int64_t topk,
95	- int64_t maxSeqlenQ, int64_t maxSeqlenK, const char layoutQOptional, const char layoutKOptional, int64_t maskMode,	95	+ int64_t batchSize, int64_t maxSeqlenQ, int64_t maxSeqlenK, const char *layoutQOptional,
96	- int64_t cmpRatio, uint32_t aicCoreNum, uint32_t aivCoreNum, const std::string &socVersion)	96	+ const char *layoutKOptional, int64_t maskMode, int64_t cmpRatio, uint32_t aicCoreNum, uint32_t aivCoreNum,
		97	+ const std::string &socVersion)
97	{	98	{
98	// num_heads_q 校验	99	// num_heads_q 校验
99	CHECK_COND(numHeadsQ >= LI_V2_NUM_HEADS_Q_LOWER_BOUND && numHeadsQ <= LI_V2_NUM_HEADS_Q_UPPER_BOUND,	100	CHECK_COND(numHeadsQ >= LI_V2_NUM_HEADS_Q_LOWER_BOUND && numHeadsQ <= LI_V2_NUM_HEADS_Q_UPPER_BOUND,
	@@ -102,6 +103,9 @@ aclnnStatus CheckSingleParamLiV2(int64_t numHeadsQ, int64_t numHeadsK, int64_t t
102	// num_heads_k 校验	103	// num_heads_k 校验
103	CHECK_COND(numHeadsK == 1, ACLNN_ERR_PARAM_INVALID,	104	CHECK_COND(numHeadsK == 1, ACLNN_ERR_PARAM_INVALID,
104	"num_heads_kv should only be 1, but got %lld", numHeadsK);	105	"num_heads_kv should only be 1, but got %lld", numHeadsK);
		106	+ // head_dim 校验
		107	+ CHECK_COND(headDim == 128, ACLNN_ERR_PARAM_INVALID,
		F F FiguraDoge8月6日魔鬼数字，可以使用set来记录允许的headdim值，这样扩充起来也方便当然直接检查也是可以的 han-dongchen 8月6日评论： 8月6日评论：
		108	+ "head_dim should be 128, but got %lld", headDim);
105	// topk 校验	109	// topk 校验
106	CHECK_COND(topk >= LI_V2_TOPK_LOWER_BOUND && topk <= LI_V2_TOPK_UPPER_BOUND, ACLNN_ERR_PARAM_INVALID,	110	CHECK_COND(topk >= LI_V2_TOPK_LOWER_BOUND && topk <= LI_V2_TOPK_UPPER_BOUND, ACLNN_ERR_PARAM_INVALID,
107	"topk should be [%lld, %lld], but got %lld", LI_V2_TOPK_LOWER_BOUND, LI_V2_TOPK_UPPER_BOUND, topk);	111	"topk should be [%lld, %lld], but got %lld", LI_V2_TOPK_LOWER_BOUND, LI_V2_TOPK_UPPER_BOUND, topk);
	@@ -136,6 +140,18 @@ aclnnStatus CheckSingleParamLiV2(int64_t numHeadsQ, int64_t numHeadsK, int64_t t
136	CHECK_COND((strcmp(layoutQOptional, layoutKOptional) == 0), ACLNN_ERR_PARAM_INVALID,	140	CHECK_COND((strcmp(layoutQOptional, layoutKOptional) == 0), ACLNN_ERR_PARAM_INVALID,
137	"For layout_k != PA_BBND, layout_q and layout_k must be the same!");	141	"For layout_k != PA_BBND, layout_q and layout_k must be the same!");
138	}	142	}
		143	+ // 校验 layout_q 为 BSND 时，max_seqlen_q 必须大于 0
		144	+ if (strcmp(layoutQOptional, "BSND") == 0) {
		145	+ CHECK_COND(maxSeqlenQ > 0, ACLNN_ERR_PARAM_INVALID,
		146	+ "When layout_q is BSND, the value of max_seqlen_q "
		147	+ "must be equal to the size of the second axis of q, but got %lld", maxSeqlenQ);
		148	+ }
		F F FiguraDoge8月6日输出日志与防御内容不匹配 han-dongchen 8月6日评论： 8月6日评论：
		149	+ // 校验 layout_k 为 BSND 时，max_seqlen_k 必须大于 0
		150	+ if (strcmp(layoutKOptional, "BSND") == 0) {
		151	+ CHECK_COND(maxSeqlenK > 0, ACLNN_ERR_PARAM_INVALID,
		152	+ "When layout_k is BSND, the value of max_seqlen_k "
		153	+ "must be equal to the size of the second axis of k, but got %lld", maxSeqlenK);
		F F FiguraDoge8月6日同一个问题，日志不匹配 han-dongchen 8月6日评论： 8月6日评论：
		154	+ }
139	// 核心数校验	155	// 核心数校验
140	CHECK_COND(aicCoreNum > 0, ACLNN_ERR_PARAM_INVALID, "AIC num should be larger than 0, but got %u", aicCoreNum);	156	CHECK_COND(aicCoreNum > 0, ACLNN_ERR_PARAM_INVALID, "AIC num should be larger than 0, but got %u", aicCoreNum);
141	CHECK_COND(aicCoreNum <= optiling::AIC_CORE_MAX_NUM, ACLNN_ERR_PARAM_INVALID,	157	CHECK_COND(aicCoreNum <= optiling::AIC_CORE_MAX_NUM, ACLNN_ERR_PARAM_INVALID,
	@@ -287,8 +303,8 @@ aclnnStatus ParamsCheckLiV2(
287	char layoutKOptional, int64_t maskMode, int64_t cmpRatio, const aclTensor metadata, uint32_t aicCoreNum,	303	char layoutKOptional, int64_t maskMode, int64_t cmpRatio, const aclTensor metadata, uint32_t aicCoreNum,
288	uint32_t aivCoreNum, const std::string &socVersion)	304	uint32_t aivCoreNum, const std::string &socVersion)
289	{	305	{
290	- auto ret = CheckSingleParamLiV2(numHeadsQ, numHeadsK, topk, batchSize, maxSeqlenQ, maxSeqlenK, layoutQOptional,	306	+ auto ret = CheckSingleParamLiV2(numHeadsQ, numHeadsK, headDim, topk, batchSize, maxSeqlenQ, maxSeqlenK,
291	- layoutKOptional, maskMode, cmpRatio, aicCoreNum, aivCoreNum, socVersion);	307	+ layoutQOptional, layoutKOptional, maskMode, cmpRatio, aicCoreNum, aivCoreNum, socVersion);
292	CHECK_RET(ret == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);	308	CHECK_RET(ret == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
293		309
294	ret = CheckExistenceLiV2(maskMode, cmpRatio, cuSeqlensQOptional, cuSeqlensKOptional, sequsedQOptional,	310	ret = CheckExistenceLiV2(maskMode, cmpRatio, cuSeqlensQOptional, cuSeqlensKOptional, sequsedQOptional,

Mattention/lightning_indexer_v2_metadata/op_kernel_aicpu/lightning_indexer_v2_metadata_aicpu.cpp+47-15

	@@ -76,13 +76,12 @@ bool LightningIndexerV2MetadataCpuKernel::ParamsCheck()
76	if (layoutQ_ == "TND") {	76	if (layoutQ_ == "TND") {
77	if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) {	77	if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) {
78	const int32_t cuSeqlensQPtr = static_cast<const int32_t>(cuSeqlensQ_->GetData());	78	const int32_t cuSeqlensQPtr = static_cast<const int32_t>(cuSeqlensQ_->GetData());
		79	+ // 校验 cu_seqlens_q 首元素为 0
		80	+ if (cuSeqlensQPtr[0] != 0) {
		81	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_q should be 0, but got %d", cuSeqlensQPtr[0]);
		82	+ return false;
		83	+ }
79	for (int i = 0; i < batchSize + 1; i++) {	84	for (int i = 0; i < batchSize + 1; i++) {
80	- // 校验 cu_seqlens_q 元素非负
81	- if (cuSeqlensQPtr[i] < 0) {
82	- KERNEL_LOG_ERROR("The elements in cu_seqlens_q should be >= 0, but got cu_seqlens_q[%d] = %d",
83	- i, cuSeqlensQPtr[i]);
84	- return false;
85	- }
86	// 校验 cu_seqlens_q 元素递增	85	// 校验 cu_seqlens_q 元素递增
87	if (i > 0 && cuSeqlensQPtr[i - 1] > cuSeqlensQPtr[i]) {	86	if (i > 0 && cuSeqlensQPtr[i - 1] > cuSeqlensQPtr[i]) {
88	KERNEL_LOG_ERROR("The elements in cu_seqlens_q must be in ascending order, "	87	KERNEL_LOG_ERROR("The elements in cu_seqlens_q must be in ascending order, "
	@@ -97,13 +96,12 @@ bool LightningIndexerV2MetadataCpuKernel::ParamsCheck()
97	if (layoutK_ == "TND") {	96	if (layoutK_ == "TND") {
98	if (cuSeqlensK_ != nullptr && cuSeqlensK_->GetData() != nullptr) {	97	if (cuSeqlensK_ != nullptr && cuSeqlensK_->GetData() != nullptr) {
99	const int32_t cuSeqlensKPtr = static_cast<const int32_t>(cuSeqlensK_->GetData());	98	const int32_t cuSeqlensKPtr = static_cast<const int32_t>(cuSeqlensK_->GetData());
		99	+ // 校验 cu_seqlens_k 首元素为 0
		100	+ if (cuSeqlensKPtr[0] != 0) {
		101	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_k should be 0, but got %d", cuSeqlensKPtr[0]);
		102	+ return false;
		103	+ }
100	for (int i = 0; i < batchSize + 1; i++) {	104	for (int i = 0; i < batchSize + 1; i++) {
101	- // 校验 cu_seqlens_k 元素非负
102	- if (cuSeqlensKPtr[i] < 0) {
103	- KERNEL_LOG_ERROR("The elements in cu_seqlens_k should be >= 0, but got cu_seqlens_k[%d] = %d",
104	- i, cuSeqlensKPtr[i]);
105	- return false;
106	- }
107	// 校验 cu_seqlens_k 元素递增	105	// 校验 cu_seqlens_k 元素递增
108	if (i > 0 && cuSeqlensKPtr[i - 1] > cuSeqlensKPtr[i]) {	106	if (i > 0 && cuSeqlensKPtr[i - 1] > cuSeqlensKPtr[i]) {
109	KERNEL_LOG_ERROR("The elements in cu_seqlens_k must be in ascending order, "	107	KERNEL_LOG_ERROR("The elements in cu_seqlens_k must be in ascending order, "
	@@ -117,32 +115,66 @@ bool LightningIndexerV2MetadataCpuKernel::ParamsCheck()
117	// 校验 seqused_q 元素非负	115	// 校验 seqused_q 元素非负
118	if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) {	116	if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) {
119	const int32_t sequsedQPtr = static_cast<const int32_t>(sequsedQ_->GetData());	117	const int32_t sequsedQPtr = static_cast<const int32_t>(sequsedQ_->GetData());
		118	+ const int32_t *cuSeqlensQPtr = (layoutQ_ == "TND" && cuSeqlensQ_ != nullptr &&
		119	+ cuSeqlensQ_->GetData() != nullptr) ?
		120	+ static_cast<const int32_t*>(cuSeqlensQ_->GetData()) : nullptr;
120	for (int i = 0; i < batchSize; i++) {	121	for (int i = 0; i < batchSize; i++) {
121	if (sequsedQPtr[i] < 0) {	122	if (sequsedQPtr[i] < 0) {
122	KERNEL_LOG_ERROR("The elements in seqused_q should be >= 0, but got seqused_q[%d] = %d",	123	KERNEL_LOG_ERROR("The elements in seqused_q should be >= 0, but got seqused_q[%d] = %d",
123	i, sequsedQPtr[i]);	124	i, sequsedQPtr[i]);
124	return false;	125	return false;
125	}	126	}
		127	+ // 校验 seqused_q 元素不大于 max_seqlen_q (BSND) 或 cu_seqlens_q 序列长度 (TND)
		128	+ if (layoutQ_ == "BSND" && sequsedQPtr[i] > maxSeqlenQ_) {
		129	+ KERNEL_LOG_ERROR("The elements in seqused_q should not be greater than max_seqlen_q %d, "
		130	+ "but got seqused_q[%d] = %d", maxSeqlenQ_, i, sequsedQPtr[i]);
		131	+ return false;
		132	+ }
		133	+ if (cuSeqlensQPtr != nullptr) {
		134	+ int32_t seqLen = cuSeqlensQPtr[i + 1] - cuSeqlensQPtr[i];
		135	+ if (sequsedQPtr[i] > seqLen) {
		136	+ KERNEL_LOG_ERROR("The elements in seqused_q should not be greater than the sequence length "
		137	+ "from cu_seqlens_q %d, but got seqused_q[%d] = %d", seqLen, i, sequsedQPtr[i]);
		138	+ return false;
		139	+ }
		140	+ }
126	}	141	}
127	}	142	}
128	// 校验 seqused_k 元素非负	143	// 校验 seqused_k 元素非负
129	if (sequsedK_ != nullptr && sequsedK_->GetData() != nullptr) {	144	if (sequsedK_ != nullptr && sequsedK_->GetData() != nullptr) {
130	const int32_t sequsedKPtr = static_cast<const int32_t>(sequsedK_->GetData());	145	const int32_t sequsedKPtr = static_cast<const int32_t>(sequsedK_->GetData());
		146	+ const int32_t *cuSeqlensKPtr = (layoutK_ == "TND" && cuSeqlensK_ != nullptr &&
		147	+ cuSeqlensK_->GetData() != nullptr) ?
		148	+ static_cast<const int32_t*>(cuSeqlensK_->GetData()) : nullptr;
131	for (int i = 0; i < batchSize; i++) {	149	for (int i = 0; i < batchSize; i++) {
132	if (sequsedKPtr[i] < 0) {	150	if (sequsedKPtr[i] < 0) {
133	KERNEL_LOG_ERROR("The elements in seqused_k should be >= 0, but got seqused_k[%d] = %d",	151	KERNEL_LOG_ERROR("The elements in seqused_k should be >= 0, but got seqused_k[%d] = %d",
134	i, sequsedKPtr[i]);	152	i, sequsedKPtr[i]);
135	return false;	153	return false;
136	}	154	}
		155	+ // 校验 seqused_k 元素不大于 max_seqlen_k (BSND) 或 cu_seqlens_k 序列长度 (TND)
		156	+ if (layoutK_ == "BSND" && sequsedKPtr[i] > maxSeqlenK_) {
		157	+ KERNEL_LOG_ERROR("The elements in seqused_k should not be greater than max_seqlen_k %d, "
		158	+ "but got seqused_k[%d] = %d", maxSeqlenK_, i, sequsedKPtr[i]);
		159	+ return false;
		160	+ }
		161	+ if (cuSeqlensKPtr != nullptr) {
		162	+ int32_t seqLen = cuSeqlensKPtr[i + 1] - cuSeqlensKPtr[i];
		163	+ if (sequsedKPtr[i] > seqLen) {
		164	+ KERNEL_LOG_ERROR("The elements in seqused_k should not be greater than the sequence length "
		165	+ "from cu_seqlens_k %d, but got seqused_k[%d] = %d", seqLen, i, sequsedKPtr[i]);
		166	+ return false;
		167	+ }
		168	+ }
137	}	169	}
138	}	170	}
139	- // 校验 cmp_residual_k 元素非负	171	+ // 校验 cmp_residual_k 元素
140	if (cmpResidualK_ != nullptr && cmpResidualK_->GetData() != nullptr) {	172	if (cmpResidualK_ != nullptr && cmpResidualK_->GetData() != nullptr) {
141	const int32_t cmpResidualKPtr = static_cast<const int32_t>(cmpResidualK_->GetData());	173	const int32_t cmpResidualKPtr = static_cast<const int32_t>(cmpResidualK_->GetData());
142	for (int i = 0; i < batchSize; i++) {	174	for (int i = 0; i < batchSize; i++) {
143	if (cmpResidualKPtr[i] < 0 \|\| cmpResidualKPtr[i] >= cmpRatio_) {	175	if (cmpResidualKPtr[i] < 0 \|\| cmpResidualKPtr[i] >= cmpRatio_) {
144	- KERNEL_LOG_ERROR("The elements in cmp_residual_k should be in [0, cmpRatio_), but got "	176	+ KERNEL_LOG_ERROR("The elements in cmp_residual_k should be in [0, cmpRatio_(%d)), but got "
145	- "cmp_residual_k[%d] = %d", i, cmpResidualKPtr[i]);	177	+ "cmp_residual_k[%d] = %d", cmpRatio_, i, cmpResidualKPtr[i]);
146	return false;	178	return false;
147	}	179	}
148	}	180	}

Mattention/mixed_quant_sparse_flash_mla_metadata/op_host/mixed_quant_sparse_flash_mla_metadata_check.h+160-3

	@@ -242,6 +242,31 @@ aclnnStatus CheckSingleParamMqsmla(int64_t batchSize, int64_t maxSeqlenQ, int64_
242	"must be equal to that of layout_kv");	242	"must be equal to that of layout_kv");
243	return ACLNN_ERR_PARAM_INVALID;	243	return ACLNN_ERR_PARAM_INVALID;
244	}	244	}
		245	+ // 校验 layout_q 为 BSND 时，max_seqlen_q 必须大于 0
		246	+ if (strcmp(layoutQOptional, "BSND") == 0 && maxSeqlenQ <= 0) {
		247	+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(MQSMLA_ACLNN_OP_NAME, "max_seqlen_q", std::to_string(maxSeqlenQ),
		248	+ "When layout_q is BSND, the value of max_seqlen_q "
		249	+ "must be equal to the size of the second axis of q");
		250	+ return ACLNN_ERR_PARAM_INVALID;
		251	+ }
		252	+ // 校验 has_ori_kv 且 layout_kv 为 BSND 时，max_seqlen_ori_kv 必须大于 0
		253	+ if (hasOriKv && strcmp(layoutKvOptional, "BSND") == 0 && maxSeqlenOriKv <= 0) {
		254	+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(MQSMLA_ACLNN_OP_NAME, "max_seqlen_ori_kv",
		255	+ std::to_string(maxSeqlenOriKv),
		256	+ "When has_ori_kv is true and layout_kv is BSND, "
		257	+ "the value of max_seqlen_ori_kv "
		258	+ "must be equal to the size of the second axis of ori_kv");
		259	+ return ACLNN_ERR_PARAM_INVALID;
		260	+ }
		261	+ // 校验 has_cmp_kv 且 layout_kv 为 BSND 时，max_seqlen_cmp_kv 必须大于 0
		262	+ if (hasCmpKv && strcmp(layoutKvOptional, "BSND") == 0 && maxSeqlenCmpKv <= 0) {
		263	+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(MQSMLA_ACLNN_OP_NAME, "max_seqlen_cmp_kv",
		264	+ std::to_string(maxSeqlenCmpKv),
		265	+ "When has_cmp_kv is true and layout_kv is BSND, "
		266	+ "the value of max_seqlen_cmp_kv "
		267	+ "must be equal to the size of the second axis of cmp_kv");
		F F FiguraDoge8月6日错误日志 han-dongchen 8月6日评论： 8月6日评论：
		268	+ return ACLNN_ERR_PARAM_INVALID;
		269	+ }
245	// 核数校验	270	// 核数校验
246	if (aicCoreNum == 0) {	271	if (aicCoreNum == 0) {
247	OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(MQSMLA_ACLNN_OP_NAME, "aic_core_num", std::to_string(aicCoreNum),	272	OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(MQSMLA_ACLNN_OP_NAME, "aic_core_num", std::to_string(aicCoreNum),
	@@ -425,13 +450,72 @@ int64_t GetCmpKvBatchSizeMqsmla(const aclTensor *sequsedCmpKvOptional, const acl
425	return batchSize;	450	return batchSize;
426	}	451	}
427		452
		453	+std::string TopkLengthShapeToStringMqsmla(const aclTensor *topkLengthOptional)
		454	+{
		455	+ const auto &shape = topkLengthOptional->GetViewShape();
		456	+ std::string result;
		457	+ for (size_t i = 0; i < shape.GetDimNum(); ++i) {
		458	+ if (i != 0) {
		459	+ result += ", ";
		460	+ }
		461	+ result += std::to_string(shape.GetDim(i));
		462	+ }
		463	+ return result;
		464	+}
		465	+
		466	+aclnnStatus CheckTopkLengthFirstDimMqsmla(const aclTensor *topkLengthOptional, const std::string &topkLengthName,
		467	+ int64_t queryBatchSize, const std::string &querySource)
		468	+{
		469	+ if (topkLengthOptional->GetViewShape().GetDim(0) == queryBatchSize) {
		470	+ return ACLNN_SUCCESS;
		471	+ }
		472	+ std::string incorrectShape = TopkLengthShapeToStringMqsmla(topkLengthOptional);
		473	+ if (IsTensorSourceMqsmla(querySource)) {
		474	+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
		475	+ MQSMLA_ACLNN_OP_NAME, topkLengthName, incorrectShape,
		476	+ "When layout_q is BSND, the size of the first axis of " + topkLengthName +
		477	+ " must be equal to " + GetSourceDescMqsmla(querySource));
		478	+ } else {
		479	+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
		480	+ MQSMLA_ACLNN_OP_NAME, topkLengthName, incorrectShape,
		481	+ "When layout_q is BSND, the size of the first axis of " + topkLengthName +
		482	+ " must be equal to batch_size");
		483	+ }
		484	+ return ACLNN_ERR_PARAM_INVALID;
		485	+}
		486	+
		487	+struct TopkLengthAxisMqsmla {
		488	+ int64_t index;
		489	+ const char *desc;
		490	+};
		491	+
		492	+inline constexpr TopkLengthAxisMqsmla MQSMLA_TOPK_LENGTH_SECOND_AXIS{1, "second"};
		493	+inline constexpr TopkLengthAxisMqsmla MQSMLA_TOPK_LENGTH_THIRD_AXIS{2, "third"};
		494	+
		495	+aclnnStatus CheckTopkLengthSingleDimMqsmla(const aclTensor *topkLengthOptional, const std::string &topkLengthName,
		496	+ TopkLengthAxisMqsmla axis, int64_t expectedValue,
		497	+ const std::string &expectedDesc, const char *layoutQOptional)
		498	+{
		499	+ if (topkLengthOptional->GetViewShape().GetDim(axis.index) == expectedValue) {
		500	+ return ACLNN_SUCCESS;
		501	+ }
		502	+ std::string incorrectShape = TopkLengthShapeToStringMqsmla(topkLengthOptional);
		503	+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
		504	+ MQSMLA_ACLNN_OP_NAME, topkLengthName, incorrectShape,
		505	+ "When layout_q is " + std::string(layoutQOptional) + ", the size of the " + axis.desc + " axis of " +
		506	+ topkLengthName + " must be equal to " + expectedDesc);
		507	+ return ACLNN_ERR_PARAM_INVALID;
		508	+}
		509	+
428	aclnnStatus CheckConsistencyMqsmla(const aclTensor cuSeqlensQOptional, const aclTensor cuSeqlensOriKvOptional,	510	aclnnStatus CheckConsistencyMqsmla(const aclTensor cuSeqlensQOptional, const aclTensor cuSeqlensOriKvOptional,
429	const aclTensor cuSeqlensCmpKvOptional, const aclTensor sequsedQOptional,	511	const aclTensor cuSeqlensCmpKvOptional, const aclTensor sequsedQOptional,
430	const aclTensor sequsedOriKvOptional, const aclTensor sequsedCmpKvOptional,	512	const aclTensor sequsedOriKvOptional, const aclTensor sequsedCmpKvOptional,
431	const aclTensor cmpResidualKvOptional, const aclTensor oriTopkLengthOptional,	513	const aclTensor cmpResidualKvOptional, const aclTensor oriTopkLengthOptional,
432	const aclTensor *cmpTopkLengthOptional, int64_t batchSize,	514	const aclTensor *cmpTopkLengthOptional, int64_t batchSize,
433	const char layoutQOptional, const char layoutKvOptional, bool hasOriKv,	515	const char layoutQOptional, const char layoutKvOptional, bool hasOriKv,
434	- bool hasCmpKv, const aclTensor *metadata)	516	+ bool hasCmpKv, int64_t oriTopk, int64_t cmpTopk, int64_t oriMaskMode,
		517	+ int64_t cmpMaskMode, int64_t maxSeqlenQ, int64_t numHeadsKv,
		518	+ const aclTensor *metadata)
435	{	519	{
436	aclDataType dataType = aclDataType::ACL_DT_UNDEFINED;	520	aclDataType dataType = aclDataType::ACL_DT_UNDEFINED;
437	int64_t dimNum = -1;	521	int64_t dimNum = -1;
	@@ -506,7 +590,8 @@ aclnnStatus CheckConsistencyMqsmla(const aclTensor *cuSeqlensQOptional, const ac
506	}	590	}
507	}	591	}
508	// 校验 ori_topk_length	592	// 校验 ori_topk_length
509	- if (IsTensorExistMqsmla(oriTopkLengthOptional)) {	593	+ if (oriTopk != 0 && oriMaskMode == static_cast<int64_t>(SparseModeMqsmla::DEFAULT_MASK) &&
		594	+ IsTensorExistMqsmla(oriTopkLengthOptional)) {
		F F FiguraDoge8月6日 aicpu common现在有个检查tensor是否存在的函数，想替换可以换了 han-dongchen 8月6日评论： 8月6日评论：
510	// 校验 ori_topk_length 维度	595	// 校验 ori_topk_length 维度
511	dimNum = GetDimNumMqsmla(oriTopkLengthOptional);	596	dimNum = GetDimNumMqsmla(oriTopkLengthOptional);
512	if (strcmp(layoutQOptional, "TND") == 0) {	597	if (strcmp(layoutQOptional, "TND") == 0) {
	@@ -593,7 +678,8 @@ aclnnStatus CheckConsistencyMqsmla(const aclTensor *cuSeqlensQOptional, const ac
593	}	678	}
594	}	679	}
595	// 校验 cmp_topk_length	680	// 校验 cmp_topk_length
596	- if (IsTensorExistMqsmla(cmpTopkLengthOptional)) {	681	+ if (cmpTopk != 0 && cmpMaskMode == static_cast<int64_t>(SparseModeMqsmla::DEFAULT_MASK) &&
		682	+ IsTensorExistMqsmla(cmpTopkLengthOptional)) {
597	// 校验 cmp_topk_length 维度	683	// 校验 cmp_topk_length 维度
598	dimNum = GetDimNumMqsmla(cmpTopkLengthOptional);	684	dimNum = GetDimNumMqsmla(cmpTopkLengthOptional);
599	if (strcmp(layoutQOptional, "TND") == 0) {	685	if (strcmp(layoutQOptional, "TND") == 0) {
	@@ -704,6 +790,41 @@ aclnnStatus CheckConsistencyMqsmla(const aclTensor *cuSeqlensQOptional, const ac
704	return ACLNN_ERR_PARAM_INVALID;	790	return ACLNN_ERR_PARAM_INVALID;
705	}	791	}
706	}	792	}
		793	+ // 校验 ori_topk_length 维度一致性
		794	+ if (oriTopk != 0 &&
		795	+ oriMaskMode == static_cast<int64_t>(SparseModeMqsmla::DEFAULT_MASK) &&
		796	+ IsTensorExistMqsmla(oriTopkLengthOptional)) {
		797	+ if (strcmp(layoutQOptional, "BSND") == 0) {
		798	+ // 校验 ori_topk_length 第一个维度
		799	+ aclnnStatus ret = CheckTopkLengthFirstDimMqsmla(oriTopkLengthOptional, "ori_topk_length",
		800	+ queryBatchSize, querySource);
		801	+ if (ret != ACLNN_SUCCESS) {
		802	+ return ret;
		803	+ }
		804	+ // 校验 ori_topk_length 第二个维度
		805	+ ret = CheckTopkLengthSingleDimMqsmla(oriTopkLengthOptional, "ori_topk_length",
		806	+ MQSMLA_TOPK_LENGTH_SECOND_AXIS, maxSeqlenQ, "max_seqlen_q",
		807	+ layoutQOptional);
		808	+ if (ret != ACLNN_SUCCESS) {
		809	+ return ret;
		810	+ }
		811	+ // 校验 ori_topk_length 第三个维度
		812	+ ret = CheckTopkLengthSingleDimMqsmla(oriTopkLengthOptional, "ori_topk_length",
		813	+ MQSMLA_TOPK_LENGTH_THIRD_AXIS, numHeadsKv, "num_heads_kv",
		814	+ layoutQOptional);
		815	+ if (ret != ACLNN_SUCCESS) {
		816	+ return ret;
		817	+ }
		818	+ } else if (strcmp(layoutQOptional, "TND") == 0) {
		819	+ // 校验 ori_topk_length 第二个维度
		820	+ aclnnStatus ret = CheckTopkLengthSingleDimMqsmla(oriTopkLengthOptional, "ori_topk_length",
		821	+ MQSMLA_TOPK_LENGTH_SECOND_AXIS, numHeadsKv,
		822	+ "num_heads_kv", layoutQOptional);
		823	+ if (ret != ACLNN_SUCCESS) {
		824	+ return ret;
		825	+ }
		826	+ }
		827	+ }
707	}	828	}
708	if (hasCmpKv) {	829	if (hasCmpKv) {
709	std::string cmpKvSource;	830	std::string cmpKvSource;
	@@ -764,6 +885,41 @@ aclnnStatus CheckConsistencyMqsmla(const aclTensor *cuSeqlensQOptional, const ac
764	return ACLNN_ERR_PARAM_INVALID;	885	return ACLNN_ERR_PARAM_INVALID;
765	}	886	}
766	}	887	}
		888	+ // 校验 cmp_topk_length 维度一致性
		889	+ if (cmpTopk != 0 &&
		890	+ cmpMaskMode == static_cast<int64_t>(SparseModeMqsmla::DEFAULT_MASK) &&
		F F FiguraDoge8月6日圈复杂度没炸么 han-dongchen 8月6日评论： 8月6日评论：
		891	+ IsTensorExistMqsmla(cmpTopkLengthOptional)) {
		892	+ if (strcmp(layoutQOptional, "BSND") == 0) {
		893	+ // 校验 cmp_topk_length 第一个维度
		894	+ aclnnStatus ret = CheckTopkLengthFirstDimMqsmla(cmpTopkLengthOptional, "cmp_topk_length",
		895	+ queryBatchSize, querySource);
		896	+ if (ret != ACLNN_SUCCESS) {
		897	+ return ret;
		898	+ }
		899	+ // 校验 cmp_topk_length 第二个维度
		900	+ ret = CheckTopkLengthSingleDimMqsmla(cmpTopkLengthOptional, "cmp_topk_length",
		901	+ MQSMLA_TOPK_LENGTH_SECOND_AXIS, maxSeqlenQ, "max_seqlen_q",
		902	+ layoutQOptional);
		903	+ if (ret != ACLNN_SUCCESS) {
		904	+ return ret;
		905	+ }
		906	+ // 校验 cmp_topk_length 第三个维度
		907	+ ret = CheckTopkLengthSingleDimMqsmla(cmpTopkLengthOptional, "cmp_topk_length",
		908	+ MQSMLA_TOPK_LENGTH_THIRD_AXIS, numHeadsKv, "num_heads_kv",
		909	+ layoutQOptional);
		910	+ if (ret != ACLNN_SUCCESS) {
		911	+ return ret;
		912	+ }
		913	+ } else if (strcmp(layoutQOptional, "TND") == 0) {
		914	+ // 校验 cmp_topk_length 第二个维度
		915	+ aclnnStatus ret = CheckTopkLengthSingleDimMqsmla(cmpTopkLengthOptional, "cmp_topk_length",
		916	+ MQSMLA_TOPK_LENGTH_SECOND_AXIS, numHeadsKv,
		917	+ "num_heads_kv", layoutQOptional);
		918	+ if (ret != ACLNN_SUCCESS) {
		919	+ return ret;
		920	+ }
		921	+ }
		922	+ }
767	}	923	}
768	return ACLNN_SUCCESS;	924	return ACLNN_SUCCESS;
769	}	925	}
	@@ -791,6 +947,7 @@ static aclnnStatus ParamsCheck(const aclTensor *cuSeqlensQOptional, const aclTen
791	CheckConsistencyMqsmla(cuSeqlensQOptional, cuSeqlensOriKvOptional, cuSeqlensCmpKvOptional, sequsedQOptional,	947	CheckConsistencyMqsmla(cuSeqlensQOptional, cuSeqlensOriKvOptional, cuSeqlensCmpKvOptional, sequsedQOptional,
792	sequsedOriKvOptional, sequsedCmpKvOptional, cmpResidualKvOptional, oriTopkLengthOptional,	948	sequsedOriKvOptional, sequsedCmpKvOptional, cmpResidualKvOptional, oriTopkLengthOptional,
793	cmpTopkLengthOptional, batchSize, layoutQOptional, layoutKvOptional, hasOriKv, hasCmpKv,	949	cmpTopkLengthOptional, batchSize, layoutQOptional, layoutKvOptional, hasOriKv, hasCmpKv,
		950	+ oriTopk, cmpTopk, oriMaskMode, cmpMaskMode, maxSeqlenQ, numHeadsKv,
794	metaData) == ACLNN_SUCCESS) {	951	metaData) == ACLNN_SUCCESS) {
795	return ACLNN_SUCCESS;	952	return ACLNN_SUCCESS;
796	} else {	953	} else {

Mattention/mixed_quant_sparse_flash_mla_metadata/op_host/op_api/aclnn_mixed_quant_sparse_flash_mla_metadata.cpp+5-3

	@@ -72,12 +72,14 @@ aclnnStatus aclnnMixedQuantSparseFlashMlaMetadataGetWorkspaceSize(
72	std::string socVersionStr = npuInfo.GetSocLongVersion();	72	std::string socVersionStr = npuInfo.GetSocLongVersion();
73	const char *socVersion = socVersionStr.c_str();	73	const char *socVersion = socVersionStr.c_str();
74		74
75	- int64_t batchConsistencyLevel = 0;	75	+ int64_t deterministicLevel = 0;
76	- aclError aclRet = aclrtGetSysParamOpt(ACL_OPT_DETERMINISTIC, &batchConsistencyLevel);	76	+ aclError aclRet = aclrtGetSysParamOpt(ACL_OPT_DETERMINISTIC, &deterministicLevel);
77	if (aclRet != ACL_SUCCESS) {	77	if (aclRet != ACL_SUCCESS) {
78	OP_LOGW("aclnnMixedQuantSparseFlashMlaMetadata unable to get system param batch consistency level.");	78	OP_LOGW("aclnnMixedQuantSparseFlashMlaMetadata unable to get system param batch consistency level.");
79	}	79	}
80	- bool isBatchConsistency = (batchConsistencyLevel == BATCH_CONSISTENCY_LEVEL);	80	+ OP_LOGD("deterministic_level=%lld", deterministicLevel);
		81	+ bool isBatchConsistency = (deterministicLevel == BATCH_CONSISTENCY_LEVEL);
		82	+
81	auto ret = ParamsCheck(cuSeqlensQOptional, cuSeqlensOriKvOptional, cuSeqlensCmpKvOptional, sequsedQOptional,	83	auto ret = ParamsCheck(cuSeqlensQOptional, cuSeqlensOriKvOptional, cuSeqlensCmpKvOptional, sequsedQOptional,
82	sequsedOriKvOptional, sequsedCmpKvOptional, cmpResidualKvOptional, oriTopkLengthOptional,	84	sequsedOriKvOptional, sequsedCmpKvOptional, cmpResidualKvOptional, oriTopkLengthOptional,
83	cmpTopkLengthOptional, numHeadsQ, numHeadsKv, headDim, quantMode, batchSize, maxSeqlenQ,	85	cmpTopkLengthOptional, numHeadsQ, numHeadsKv, headDim, quantMode, batchSize, maxSeqlenQ,

Mattention/mixed_quant_sparse_flash_mla_metadata/op_kernel_aicpu/mixed_quant_sparse_flash_mla_metadata_aicpu.cpp+88-31

	@@ -92,13 +92,12 @@ bool MixedQuantSparseFlashMlaMetadataCpuKernel::ParamsCheck()
92	if (layoutQ_ == "TND") {	92	if (layoutQ_ == "TND") {
93	if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) {	93	if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) {
94	const int32_t cuSeqlensQPtr = static_cast<const int32_t >(cuSeqlensQ_->GetData());	94	const int32_t cuSeqlensQPtr = static_cast<const int32_t >(cuSeqlensQ_->GetData());
		95	+ // 校验 cu_seqlens_q 首元素为 0
		96	+ if (cuSeqlensQPtr[0] != 0) {
		97	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_q should be 0, but got %d", cuSeqlensQPtr[0]);
		98	+ return false;
		99	+ }
95	for (int i = 0; i < batchSize + 1; i++) {	100	for (int i = 0; i < batchSize + 1; i++) {
96	- // 校验 cu_seqlens_q 元素非负
97	- if (cuSeqlensQPtr[i] < 0) {
98	- KERNEL_LOG_ERROR("The elements in cu_seqlens_q should be >= 0, but got cu_seqlens_q[%d] = %d", i,
99	- cuSeqlensQPtr[i]);
100	- return false;
101	- }
102	// 校验 cu_seqlens_q 元素递增	101	// 校验 cu_seqlens_q 元素递增
103	if (i > 0 && cuSeqlensQPtr[i - 1] > cuSeqlensQPtr[i]) {	102	if (i > 0 && cuSeqlensQPtr[i - 1] > cuSeqlensQPtr[i]) {
104	KERNEL_LOG_ERROR("The elements in cu_seqlens_q must be in ascending order, "	103	KERNEL_LOG_ERROR("The elements in cu_seqlens_q must be in ascending order, "
	@@ -112,6 +111,9 @@ bool MixedQuantSparseFlashMlaMetadataCpuKernel::ParamsCheck()
112	// 校验 seqused_q 元素	111	// 校验 seqused_q 元素
113	if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) {	112	if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) {
114	const int32_t sequsedQPtr = static_cast<const int32_t >(sequsedQ_->GetData());	113	const int32_t sequsedQPtr = static_cast<const int32_t >(sequsedQ_->GetData());
		114	+ const int32_t *cuSeqlensQPtr = (layoutQ_ == "TND" && cuSeqlensQ_ != nullptr &&
		115	+ cuSeqlensQ_->GetData() != nullptr) ?
		F F FiguraDoge8月6日我理解直接cuSeqlensQ->GetData就行了，这个三元运算符的意义在哪里 han-dongchen 8月6日评论： 8月6日评论：
		116	+ static_cast<const int32_t *>(cuSeqlensQ_->GetData()) : nullptr;
115	for (int i = 0; i < batchSize; i++) {	117	for (int i = 0; i < batchSize; i++) {
116	// 校验 seqused_q 元素非负	118	// 校验 seqused_q 元素非负
117	if (sequsedQPtr[i] < 0) {	119	if (sequsedQPtr[i] < 0) {
	@@ -119,6 +121,20 @@ bool MixedQuantSparseFlashMlaMetadataCpuKernel::ParamsCheck()
119	sequsedQPtr[i]);	121	sequsedQPtr[i]);
120	return false;	122	return false;
121	}	123	}
		124	+ // 校验 seqused_q 元素不大于 max_seqlen_q (BSND) 或 cu_seqlens_q 序列长度 (TND)
		125	+ if (layoutQ_ == "BSND" && sequsedQPtr[i] > maxSeqlenQ_) {
		126	+ KERNEL_LOG_ERROR("The elements in seqused_q should not be greater than max_seqlen_q %d, "
		127	+ "but got seqused_q[%d] = %d", maxSeqlenQ_, i, sequsedQPtr[i]);
		128	+ return false;
		129	+ }
		130	+ if (cuSeqlensQPtr != nullptr) {
		131	+ int32_t seqLen = cuSeqlensQPtr[i + 1] - cuSeqlensQPtr[i];
		132	+ if (sequsedQPtr[i] > seqLen) {
		133	+ KERNEL_LOG_ERROR("The elements in seqused_q should not be greater than the sequence length "
		134	+ "from cu_seqlens_q %d, but got seqused_q[%d] = %d", seqLen, i, sequsedQPtr[i]);
		135	+ return false;
		136	+ }
		137	+ }
122	}	138	}
123	}	139	}
124	if (hasOriKv_) {	140	if (hasOriKv_) {
	@@ -126,14 +142,13 @@ bool MixedQuantSparseFlashMlaMetadataCpuKernel::ParamsCheck()
126	if (layoutKv_ == "TND") {	142	if (layoutKv_ == "TND") {
127	if (cuSeqlensOriKv_ != nullptr && cuSeqlensOriKv_->GetData() != nullptr) {	143	if (cuSeqlensOriKv_ != nullptr && cuSeqlensOriKv_->GetData() != nullptr) {
128	const int32_t cuSeqlensOriKvPtr = static_cast<const int32_t >(cuSeqlensOriKv_->GetData());	144	const int32_t cuSeqlensOriKvPtr = static_cast<const int32_t >(cuSeqlensOriKv_->GetData());
		145	+ // 校验 cu_seqlens_ori_kv 首元素为 0
		146	+ if (cuSeqlensOriKvPtr[0] != 0) {
		147	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_ori_kv should be 0, but got %d",
		148	+ cuSeqlensOriKvPtr[0]);
		149	+ return false;
		150	+ }
129	for (int i = 0; i < batchSize + 1; i++) {	151	for (int i = 0; i < batchSize + 1; i++) {
130	- // 校验 cu_seqlens_ori_kv 元素非负
131	- if (cuSeqlensOriKvPtr[i] < 0) {
132	- KERNEL_LOG_ERROR("The elements in cu_seqlens_ori_kv should be >= 0, "
133	- "but got cu_seqlens_ori_kv[%d] = %d",
134	- i, cuSeqlensOriKvPtr[i]);
135	- return false;
136	- }
137	// 校验 cu_seqlens_ori_kv 元素递增	152	// 校验 cu_seqlens_ori_kv 元素递增
138	if (i > 0 && cuSeqlensOriKvPtr[i - 1] > cuSeqlensOriKvPtr[i]) {	153	if (i > 0 && cuSeqlensOriKvPtr[i - 1] > cuSeqlensOriKvPtr[i]) {
139	KERNEL_LOG_ERROR("The elements in cu_seqlens_ori_kv must be in ascending order, "	154	KERNEL_LOG_ERROR("The elements in cu_seqlens_ori_kv must be in ascending order, "
	@@ -147,6 +162,10 @@ bool MixedQuantSparseFlashMlaMetadataCpuKernel::ParamsCheck()
147	// 校验 seqused_ori_kv 元素	162	// 校验 seqused_ori_kv 元素
148	if (sequsedOriKv_ != nullptr && sequsedOriKv_->GetData() != nullptr) {	163	if (sequsedOriKv_ != nullptr && sequsedOriKv_->GetData() != nullptr) {
149	const int32_t sequsedOriKvPtr = static_cast<const int32_t >(sequsedOriKv_->GetData());	164	const int32_t sequsedOriKvPtr = static_cast<const int32_t >(sequsedOriKv_->GetData());
		165	+ const int32_t *cuSeqlensOriKvPtr = (layoutKv_ == "TND" && cuSeqlensOriKv_ != nullptr &&
		166	+ cuSeqlensOriKv_->GetData() != nullptr) ?
		167	+ static_cast<const int32_t *>(cuSeqlensOriKv_->GetData()) :
		168	+ nullptr;
150	for (int i = 0; i < batchSize; i++) {	169	for (int i = 0; i < batchSize; i++) {
151	// 校验 seqused_ori_kv 元素非负	170	// 校验 seqused_ori_kv 元素非负
152	if (sequsedOriKvPtr[i] < 0) {	171	if (sequsedOriKvPtr[i] < 0) {
	@@ -154,10 +173,27 @@ bool MixedQuantSparseFlashMlaMetadataCpuKernel::ParamsCheck()
154	i, sequsedOriKvPtr[i]);	173	i, sequsedOriKvPtr[i]);
155	return false;	174	return false;
156	}	175	}
		176	+ // 校验 seqused_ori_kv 元素不大于 max_seqlen_ori_kv (BSND) 或 cu_seqlens_ori_kv 序列长度 (TND)
		177	+ if (layoutKv_ == "BSND" && sequsedOriKvPtr[i] > maxSeqlenOriKv_) {
		178	+ KERNEL_LOG_ERROR("The elements in seqused_ori_kv should not be greater than "
		179	+ "max_seqlen_ori_kv %d, but got seqused_ori_kv[%d] = %d",
		180	+ maxSeqlenOriKv_, i, sequsedOriKvPtr[i]);
		181	+ return false;
		182	+ }
		183	+ if (cuSeqlensOriKvPtr != nullptr) {
		184	+ int32_t seqLen = cuSeqlensOriKvPtr[i + 1] - cuSeqlensOriKvPtr[i];
		185	+ if (sequsedOriKvPtr[i] > seqLen) {
		186	+ KERNEL_LOG_ERROR("The elements in seqused_ori_kv should not be greater than the sequence "
		187	+ "length from cu_seqlens_ori_kv %d, but got seqused_ori_kv[%d] = %d",
		188	+ seqLen, i, sequsedOriKvPtr[i]);
		189	+ return false;
		190	+ }
		191	+ }
157	}	192	}
158	}	193	}
159	// 校验 ori_topk_length 元素	194	// 校验 ori_topk_length 元素
160	- if (oriTopkLength_ != nullptr && oriTopkLength_->GetData() != nullptr) {	195	+ if (oriTopK_ != 0 && oriMaskMode_ == static_cast<int32_t>(SparseMode::DEFAULT_MASK) &&
		196	+ oriTopkLength_ != nullptr && oriTopkLength_->GetData() != nullptr) {
161	// 校验 ori_topk_length 元素数量	197	// 校验 ori_topk_length 元素数量
162	int32_t sumOfQuerySeq = GetSumOfQuerySeq();	198	int32_t sumOfQuerySeq = GetSumOfQuerySeq();
163	const int32_t oriTopkLengthPtr = static_cast<const int32_t >(oriTopkLength_->GetData());	199	const int32_t oriTopkLengthPtr = static_cast<const int32_t >(oriTopkLength_->GetData());
	@@ -187,14 +223,13 @@ bool MixedQuantSparseFlashMlaMetadataCpuKernel::ParamsCheck()
187	// 校验 cu_seqlens_cmp_kv 元素	223	// 校验 cu_seqlens_cmp_kv 元素
188	if (cuSeqlensCmpKv_ != nullptr && cuSeqlensCmpKv_->GetData() != nullptr) {	224	if (cuSeqlensCmpKv_ != nullptr && cuSeqlensCmpKv_->GetData() != nullptr) {
189	const int32_t cuSeqlensCmpKvPtr = static_cast<const int32_t >(cuSeqlensCmpKv_->GetData());	225	const int32_t cuSeqlensCmpKvPtr = static_cast<const int32_t >(cuSeqlensCmpKv_->GetData());
		226	+ // 校验 cu_seqlens_cmp_kv 首元素为 0
		227	+ if (cuSeqlensCmpKvPtr[0] != 0) {
		228	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_cmp_kv should be 0, but got %d",
		229	+ cuSeqlensCmpKvPtr[0]);
		230	+ return false;
		231	+ }
190	for (int i = 0; i < batchSize + 1; i++) {	232	for (int i = 0; i < batchSize + 1; i++) {
191	- // 校验 cu_seqlens_cmp_kv 元素非负
192	- if (cuSeqlensCmpKvPtr[i] < 0) {
193	- KERNEL_LOG_ERROR("The elements in cu_seqlens_cmp_kv should be >= 0, "
194	- "but got cu_seqlens_cmp_kv[%d] = %d",
195	- i, cuSeqlensCmpKvPtr[i]);
196	- return false;
197	- }
198	// 校验 cu_seqlens_cmp_kv 元素递增	233	// 校验 cu_seqlens_cmp_kv 元素递增
199	if (i > 0 && cuSeqlensCmpKvPtr[i - 1] > cuSeqlensCmpKvPtr[i]) {	234	if (i > 0 && cuSeqlensCmpKvPtr[i - 1] > cuSeqlensCmpKvPtr[i]) {
200	KERNEL_LOG_ERROR("The elements in cu_seqlens_cmp_kv must be in ascending order, "	235	KERNEL_LOG_ERROR("The elements in cu_seqlens_cmp_kv must be in ascending order, "
	@@ -208,6 +243,10 @@ bool MixedQuantSparseFlashMlaMetadataCpuKernel::ParamsCheck()
208	// 校验 seqused_cmp_kv 元素	243	// 校验 seqused_cmp_kv 元素
209	if (sequsedCmpKv_ != nullptr && sequsedCmpKv_->GetData() != nullptr) {	244	if (sequsedCmpKv_ != nullptr && sequsedCmpKv_->GetData() != nullptr) {
210	const int32_t sequsedCmpKvPtr = static_cast<const int32_t >(sequsedCmpKv_->GetData());	245	const int32_t sequsedCmpKvPtr = static_cast<const int32_t >(sequsedCmpKv_->GetData());
		246	+ const int32_t *cuSeqlensCmpKvPtr = (layoutKv_ == "TND" && cuSeqlensCmpKv_ != nullptr &&
		247	+ cuSeqlensCmpKv_->GetData() != nullptr) ?
		248	+ static_cast<const int32_t *>(cuSeqlensCmpKv_->GetData()) :
		249	+ nullptr;
211	for (int i = 0; i < batchSize; i++) {	250	for (int i = 0; i < batchSize; i++) {
212	// 校验 seqused_cmp_kv 元素非负	251	// 校验 seqused_cmp_kv 元素非负
213	if (sequsedCmpKvPtr[i] < 0) {	252	if (sequsedCmpKvPtr[i] < 0) {
	@@ -215,23 +254,39 @@ bool MixedQuantSparseFlashMlaMetadataCpuKernel::ParamsCheck()
215	i, sequsedCmpKvPtr[i]);	254	i, sequsedCmpKvPtr[i]);
216	return false;	255	return false;
217	}	256	}
		257	+ // 校验 seqused_cmp_kv 元素不大于 max_seqlen_cmp_kv (BSND) 或 cu_seqlens_cmp_kv 序列长度 (TND)
		258	+ if (layoutKv_ == "BSND" && sequsedCmpKvPtr[i] > maxSeqlenCmpKv_) {
		259	+ KERNEL_LOG_ERROR("The elements in seqused_cmp_kv should not be greater than "
		260	+ "max_seqlen_cmp_kv %d, but got seqused_cmp_kv[%d] = %d",
		261	+ maxSeqlenCmpKv_, i, sequsedCmpKvPtr[i]);
		262	+ return false;
		263	+ }
		264	+ if (cuSeqlensCmpKvPtr != nullptr) {
		265	+ int32_t seqLen = cuSeqlensCmpKvPtr[i + 1] - cuSeqlensCmpKvPtr[i];
		266	+ if (sequsedCmpKvPtr[i] > seqLen) {
		267	+ KERNEL_LOG_ERROR("The elements in seqused_cmp_kv should not be greater than the sequence "
		268	+ "length from cu_seqlens_cmp_kv %d, but got seqused_cmp_kv[%d] = %d",
		269	+ seqLen, i, sequsedCmpKvPtr[i]);
		270	+ return false;
		271	+ }
		272	+ }
218	}	273	}
219	}	274	}
220	// 校验 cmp_residual_kv 元素	275	// 校验 cmp_residual_kv 元素
221	if (cmpResidualKv_ != nullptr && cmpResidualKv_->GetData() != nullptr) {	276	if (cmpResidualKv_ != nullptr && cmpResidualKv_->GetData() != nullptr) {
222	const int32_t cmpResidualKvPtr = static_cast<const int32_t >(cmpResidualKv_->GetData());	277	const int32_t cmpResidualKvPtr = static_cast<const int32_t >(cmpResidualKv_->GetData());
223	for (int i = 0; i < batchSize; i++) {	278	for (int i = 0; i < batchSize; i++) {
224	- // 校验 cmp_residual_kv 元素非负
225	if (cmpResidualKvPtr[i] < 0 \|\| cmpResidualKvPtr[i] >= cmpRatio_) {	279	if (cmpResidualKvPtr[i] < 0 \|\| cmpResidualKvPtr[i] >= cmpRatio_) {
226	- KERNEL_LOG_ERROR("The elements in cmp_residual_kv should be in [0, cmpRatio_), but got "	280	+ KERNEL_LOG_ERROR("The elements in cmp_residual_kv should be in [0, cmpRatio_(%d)), but got "
227	- "cmp_residual_kv[%d] = %d",	281	+ "cmp_residual_kv[%d] = %d", cmpRatio_,
228	i, cmpResidualKvPtr[i]);	282	i, cmpResidualKvPtr[i]);
229	return false;	283	return false;
230	}	284	}
231	}	285	}
232	}	286	}
233	// 校验 cmp_topk_length 元素	287	// 校验 cmp_topk_length 元素
234	- if (cmpTopkLength_ != nullptr && cmpTopkLength_->GetData() != nullptr) {	288	+ if (cmpTopK_ != 0 && cmpMaskMode_ == static_cast<int32_t>(SparseMode::DEFAULT_MASK) &&
		289	+ cmpTopkLength_ != nullptr && cmpTopkLength_->GetData() != nullptr) {
235	// 校验 cmp_topk_length 元素数量	290	// 校验 cmp_topk_length 元素数量
236	int32_t sumOfQuerySeq = GetSumOfQuerySeq();	291	int32_t sumOfQuerySeq = GetSumOfQuerySeq();
237	const int32_t cmpTopkLengthPtr = static_cast<const int32_t >(cmpTopkLength_->GetData());	292	const int32_t cmpTopkLengthPtr = static_cast<const int32_t >(cmpTopkLength_->GetData());
	@@ -412,7 +467,8 @@ uint32_t MixedQuantSparseFlashMlaMetadataCpuKernel::GetBsStride(uint32_t bIdx, u
412	uint32_t MixedQuantSparseFlashMlaMetadataCpuKernel::GetOriTopkLength(uint32_t bsStride)	467	uint32_t MixedQuantSparseFlashMlaMetadataCpuKernel::GetOriTopkLength(uint32_t bsStride)
413	{	468	{
414	// 尝试使用 oriTopkLength_	469	// 尝试使用 oriTopkLength_
415	- if (oriTopkLength_ != nullptr && oriTopkLength_->GetData() != nullptr) {	470	+ if (oriTopK_ != 0 && oriMaskMode_ == static_cast<int32_t>(SparseMode::DEFAULT_MASK) &&
		471	+ oriTopkLength_ != nullptr && oriTopkLength_->GetData() != nullptr) {
416	const int32_t oriTopkPtr = static_cast<const int32_t >(oriTopkLength_->GetData());	472	const int32_t oriTopkPtr = static_cast<const int32_t >(oriTopkLength_->GetData());
417	return static_cast<uint32_t>(oriTopkPtr[bsStride]);	473	return static_cast<uint32_t>(oriTopkPtr[bsStride]);
418	}	474	}
	@@ -423,7 +479,8 @@ uint32_t MixedQuantSparseFlashMlaMetadataCpuKernel::GetOriTopkLength(uint32_t bs
423	uint32_t MixedQuantSparseFlashMlaMetadataCpuKernel::GetCmpTopkLength(uint32_t bsStride)	479	uint32_t MixedQuantSparseFlashMlaMetadataCpuKernel::GetCmpTopkLength(uint32_t bsStride)
424	{	480	{
425	// 尝试使用 cmpTopkLength_	481	// 尝试使用 cmpTopkLength_
426	- if (cmpTopkLength_ != nullptr && cmpTopkLength_->GetData() != nullptr) {	482	+ if (cmpTopK_ != 0 && cmpMaskMode_ == static_cast<int32_t>(SparseMode::DEFAULT_MASK) &&
		483	+ cmpTopkLength_ != nullptr && cmpTopkLength_->GetData() != nullptr) {
427	const int32_t cmpTopkPtr = static_cast<const int32_t >(cmpTopkLength_->GetData());	484	const int32_t cmpTopkPtr = static_cast<const int32_t >(cmpTopkLength_->GetData());
428	return static_cast<uint32_t>(cmpTopkPtr[bsStride]);	485	return static_cast<uint32_t>(cmpTopkPtr[bsStride]);
429	}	486	}
	@@ -465,8 +522,8 @@ uint32_t MixedQuantSparseFlashMlaMetadataCpuKernel::GetOriS2SeqSize(uint32_t bId
465	return static_cast<uint32_t>(s2Ptr[bIdx + 1U] - s2Ptr[bIdx]);	522	return static_cast<uint32_t>(s2Ptr[bIdx + 1U] - s2Ptr[bIdx]);
466	}	523	}
467	}	524	}
468	- // 如果 max_seqlen_ori_kv 没传入，且 ori_kv 为稀疏的，则尝试从 topk 中获取	525	+ // 如果是PA场景，或 max_seqlen_ori_kv 没传入，且 ori_kv 为稀疏的，则尝试从 topk 中获取
469	- if (maxSeqlenOriKv_ == 0 && isSparseOriKv_) {	526	+ if ((layoutKv_ == "PA_BBND" \|\| maxSeqlenOriKv_ == 0) && isSparseOriKv_) {
470	return UINT32_MAX;	527	return UINT32_MAX;
471	}	528	}
472	// 使用 max_seqlen_ori_kv	529	// 使用 max_seqlen_ori_kv
	@@ -488,8 +545,8 @@ uint32_t MixedQuantSparseFlashMlaMetadataCpuKernel::GetCmpS2SeqSize(uint32_t bId
488	return static_cast<uint32_t>(s2Ptr[bIdx + 1U] - s2Ptr[bIdx]);	545	return static_cast<uint32_t>(s2Ptr[bIdx + 1U] - s2Ptr[bIdx]);
489	}	546	}
490	}	547	}
491	- // 如果 max_seqlen_cmp_kv 没传入，且 cmp_kv 为稀疏的，则尝试从topk中获取	548	+ // 如果是PA场景，或 max_seqlen_cmp_kv 没传入，且 cmp_kv 为稀疏的，则尝试从topk中获取
492	- if (maxSeqlenCmpKv_ == 0 && isSparseCmpKv_) {	549	+ if ((layoutKv_ == "PA_BBND" \|\| maxSeqlenCmpKv_ == 0) && isSparseCmpKv_) {
493	return UINT32_MAX;	550	return UINT32_MAX;
494	}	551	}
495	// 使用 max_seqlen_cmp_kv	552	// 使用 max_seqlen_cmp_kv

Mattention/quant_lightning_indexer_v2_metadata/op_host/quant_lightning_indexer_v2_metadata_check.h+19-4

	@@ -99,14 +99,17 @@ int64_t GetKeyBatchSizeQliV2(int64_t batchSize, const aclTensor *cuSeqlensKOptio
99	return batchSize;	99	return batchSize;
100	}	100	}
101		101
102	-aclnnStatus CheckSingleParamQliV2(int64_t numHeadsQ, int64_t numHeadsK, int64_t topk, int64_t quantMode,	102	+aclnnStatus CheckSingleParamQliV2(int64_t numHeadsQ, int64_t numHeadsK, int64_t headDim, int64_t topk,
103	- int64_t batchSize, int64_t maxSeqlenQ, int64_t maxSeqlenK,	103	+ int64_t quantMode, int64_t batchSize, int64_t maxSeqlenQ, int64_t maxSeqlenK,
104	const char layoutQOptional, const char layoutKOptional, int64_t maskMode,	104	const char layoutQOptional, const char layoutKOptional, int64_t maskMode,
105	int64_t cmpRatio, uint32_t aicCoreNum, uint32_t aivCoreNum,	105	int64_t cmpRatio, uint32_t aicCoreNum, uint32_t aivCoreNum,
106	const std::string &socVersion)	106	const std::string &socVersion)
107	{	107	{
108	// num_heads_k 校验	108	// num_heads_k 校验
109	CHECK_COND(numHeadsK == 1, ACLNN_ERR_PARAM_INVALID, "num_heads_kv should only be 1, but got %lld", numHeadsK);	109	CHECK_COND(numHeadsK == 1, ACLNN_ERR_PARAM_INVALID, "num_heads_kv should only be 1, but got %lld", numHeadsK);
		110	+ // head_dim 校验
		111	+ CHECK_COND(headDim == 128, ACLNN_ERR_PARAM_INVALID,
		112	+ "head_dim should be 128, but got %lld", headDim);
110	// batch_size 非负校验	113	// batch_size 非负校验
111	CHECK_COND(batchSize >= 0, ACLNN_ERR_PARAM_INVALID, "batch_size should not be negative, but got %lld", batchSize);	114	CHECK_COND(batchSize >= 0, ACLNN_ERR_PARAM_INVALID, "batch_size should not be negative, but got %lld", batchSize);
112	// max_seqlen_q 校验	115	// max_seqlen_q 校验
	@@ -175,6 +178,18 @@ aclnnStatus CheckSingleParamQliV2(int64_t numHeadsQ, int64_t numHeadsK, int64_t
175	ACLNN_ERR_PARAM_INVALID,	178	ACLNN_ERR_PARAM_INVALID,
176	"layout_k must be TND/BSND/PA_BBND, but got %s", layoutKOptional);	179	"layout_k must be TND/BSND/PA_BBND, but got %s", layoutKOptional);
177	}	180	}
		181	+ // 校验 layout_q 为 BSND 时，max_seqlen_q 必须大于 0
		182	+ if (strcmp(layoutQOptional, "BSND") == 0) {
		183	+ CHECK_COND(maxSeqlenQ > 0, ACLNN_ERR_PARAM_INVALID,
		184	+ "When layout_q is BSND, the value of max_seqlen_q "
		185	+ "must be equal to the size of the second axis of q, but got %lld", maxSeqlenQ);
		186	+ }
		187	+ // 校验 layout_k 为 BSND 时，max_seqlen_k 必须大于 0
		188	+ if (strcmp(layoutKOptional, "BSND") == 0) {
		189	+ CHECK_COND(maxSeqlenK > 0, ACLNN_ERR_PARAM_INVALID,
		190	+ "When layout_k is BSND, the value of max_seqlen_k "
		191	+ "must be equal to the size of the second axis of k, but got %lld", maxSeqlenK);
		192	+ }
178	// 核心数校验	193	// 核心数校验
179	CHECK_COND(aicCoreNum > 0, ACLNN_ERR_PARAM_INVALID, "AIC num should be larger than 0, but got %u", aicCoreNum);	194	CHECK_COND(aicCoreNum > 0, ACLNN_ERR_PARAM_INVALID, "AIC num should be larger than 0, but got %u", aicCoreNum);
180	CHECK_COND(aicCoreNum <= optiling::AIC_CORE_MAX_NUM, ACLNN_ERR_PARAM_INVALID,	195	CHECK_COND(aicCoreNum <= optiling::AIC_CORE_MAX_NUM, ACLNN_ERR_PARAM_INVALID,
	@@ -332,8 +347,8 @@ aclnnStatus ParamsCheckQliV2(const aclTensor *cuSeqlensQOptional, const aclTenso
332	const std::string &socVersion)	347	const std::string &socVersion)
333	{	348	{
334	auto ret =	349	auto ret =
335	- CheckSingleParamQliV2(numHeadsQ, numHeadsK, topk, quantMode, batchSize, maxSeqlenQ, maxSeqlenK, layoutQOptional,	350	+ CheckSingleParamQliV2(numHeadsQ, numHeadsK, headDim, topk, quantMode, batchSize, maxSeqlenQ, maxSeqlenK,
336	- layoutKOptional, maskMode, cmpRatio, aicCoreNum, aivCoreNum, socVersion);	351	+ layoutQOptional, layoutKOptional, maskMode, cmpRatio, aicCoreNum, aivCoreNum, socVersion);
337	CHECK_RET(ret == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);	352	CHECK_RET(ret == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
338		353
339	ret = CheckExistenceQliV2(maskMode, cmpRatio, cuSeqlensQOptional, cuSeqlensKOptional, sequsedQOptional,	354	ret = CheckExistenceQliV2(maskMode, cmpRatio, cuSeqlensQOptional, cuSeqlensKOptional, sequsedQOptional,

Mattention/quant_lightning_indexer_v2_metadata/op_kernel_aicpu/quant_lightning_indexer_v2_metadata_aicpu.cpp+47-15

	@@ -74,13 +74,12 @@ bool QuantLightningIndexerV2MetadataCpuKernel::ParamsCheck()
74	if (layoutQ_ == "TND") {	74	if (layoutQ_ == "TND") {
75	if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) {	75	if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) {
76	const int32_t cuSeqlensQPtr = static_cast<const int32_t >(cuSeqlensQ_->GetData());	76	const int32_t cuSeqlensQPtr = static_cast<const int32_t >(cuSeqlensQ_->GetData());
		77	+ // 校验 cu_seqlens_q 首元素为 0
		78	+ if (cuSeqlensQPtr[0] != 0) {
		79	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_q should be 0, but got %d", cuSeqlensQPtr[0]);
		80	+ return false;
		81	+ }
77	for (int i = 0; i < batchSize + 1; i++) {	82	for (int i = 0; i < batchSize + 1; i++) {
78	- // 校验 cu_seqlens_q 元素非负
79	- if (cuSeqlensQPtr[i] < 0) {
80	- KERNEL_LOG_ERROR("The elements in cu_seqlens_q should be >= 0, but got cu_seqlens_q[%d] = %d", i,
81	- cuSeqlensQPtr[i]);
82	- return false;
83	- }
84	// 校验 cu_seqlens_q 元素递增	83	// 校验 cu_seqlens_q 元素递增
85	if (i > 0 && cuSeqlensQPtr[i - 1] > cuSeqlensQPtr[i]) {	84	if (i > 0 && cuSeqlensQPtr[i - 1] > cuSeqlensQPtr[i]) {
86	KERNEL_LOG_ERROR("The elements in cu_seqlens_q must be in ascending order, "	85	KERNEL_LOG_ERROR("The elements in cu_seqlens_q must be in ascending order, "
	@@ -95,13 +94,12 @@ bool QuantLightningIndexerV2MetadataCpuKernel::ParamsCheck()
95	if (layoutK_ == "TND") {	94	if (layoutK_ == "TND") {
96	if (cuSeqlensK_ != nullptr && cuSeqlensK_->GetData() != nullptr) {	95	if (cuSeqlensK_ != nullptr && cuSeqlensK_->GetData() != nullptr) {
97	const int32_t cuSeqlensKPtr = static_cast<const int32_t >(cuSeqlensK_->GetData());	96	const int32_t cuSeqlensKPtr = static_cast<const int32_t >(cuSeqlensK_->GetData());
		97	+ // 校验 cu_seqlens_k 首元素为 0
		98	+ if (cuSeqlensKPtr[0] != 0) {
		99	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_k should be 0, but got %d", cuSeqlensKPtr[0]);
		100	+ return false;
		101	+ }
98	for (int i = 0; i < batchSize + 1; i++) {	102	for (int i = 0; i < batchSize + 1; i++) {
99	- // 校验 cu_seqlens_k 元素非负
100	- if (cuSeqlensKPtr[i] < 0) {
101	- KERNEL_LOG_ERROR("The elements in cu_seqlens_k should be >= 0, but got cu_seqlens_k[%d] = %d", i,
102	- cuSeqlensKPtr[i]);
103	- return false;
104	- }
105	// 校验 cu_seqlens_k 元素递增	103	// 校验 cu_seqlens_k 元素递增
106	if (i > 0 && cuSeqlensKPtr[i - 1] > cuSeqlensKPtr[i]) {	104	if (i > 0 && cuSeqlensKPtr[i - 1] > cuSeqlensKPtr[i]) {
107	KERNEL_LOG_ERROR("The elements in cu_seqlens_k must be in ascending order, "	105	KERNEL_LOG_ERROR("The elements in cu_seqlens_k must be in ascending order, "
	@@ -115,32 +113,66 @@ bool QuantLightningIndexerV2MetadataCpuKernel::ParamsCheck()
115	// 校验 seqused_q 元素非负	113	// 校验 seqused_q 元素非负
116	if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) {	114	if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) {
117	const int32_t sequsedQPtr = static_cast<const int32_t >(sequsedQ_->GetData());	115	const int32_t sequsedQPtr = static_cast<const int32_t >(sequsedQ_->GetData());
		116	+ const int32_t *cuSeqlensQPtr = (layoutQ_ == "TND" && cuSeqlensQ_ != nullptr &&
		117	+ cuSeqlensQ_->GetData() != nullptr) ?
		118	+ static_cast<const int32_t *>(cuSeqlensQ_->GetData()) : nullptr;
118	for (int i = 0; i < batchSize; i++) {	119	for (int i = 0; i < batchSize; i++) {
119	if (sequsedQPtr[i] < 0) {	120	if (sequsedQPtr[i] < 0) {
120	KERNEL_LOG_ERROR("The elements in seqused_q should be >= 0, but got seqused_q[%d] = %d", i,	121	KERNEL_LOG_ERROR("The elements in seqused_q should be >= 0, but got seqused_q[%d] = %d", i,
121	sequsedQPtr[i]);	122	sequsedQPtr[i]);
122	return false;	123	return false;
123	}	124	}
		125	+ // 校验 seqused_q 元素不大于 max_seqlen_q (BSND) 或 cu_seqlens_q 序列长度 (TND)
		126	+ if (layoutQ_ == "BSND" && sequsedQPtr[i] > maxSeqlenQ_) {
		127	+ KERNEL_LOG_ERROR("The elements in seqused_q should not be greater than max_seqlen_q %d, "
		128	+ "but got seqused_q[%d] = %d", maxSeqlenQ_, i, sequsedQPtr[i]);
		129	+ return false;
		130	+ }
		131	+ if (cuSeqlensQPtr != nullptr) {
		132	+ int32_t seqLen = cuSeqlensQPtr[i + 1] - cuSeqlensQPtr[i];
		133	+ if (sequsedQPtr[i] > seqLen) {
		134	+ KERNEL_LOG_ERROR("The elements in seqused_q should not be greater than the sequence length "
		135	+ "from cu_seqlens_q %d, but got seqused_q[%d] = %d", seqLen, i, sequsedQPtr[i]);
		136	+ return false;
		137	+ }
		138	+ }
124	}	139	}
125	}	140	}
126	// 校验 seqused_k 元素非负	141	// 校验 seqused_k 元素非负
127	if (sequsedK_ != nullptr && sequsedK_->GetData() != nullptr) {	142	if (sequsedK_ != nullptr && sequsedK_->GetData() != nullptr) {
128	const int32_t sequsedKPtr = static_cast<const int32_t >(sequsedK_->GetData());	143	const int32_t sequsedKPtr = static_cast<const int32_t >(sequsedK_->GetData());
		144	+ const int32_t *cuSeqlensKPtr = (layoutK_ == "TND" && cuSeqlensK_ != nullptr &&
		145	+ cuSeqlensK_->GetData() != nullptr) ?
		146	+ static_cast<const int32_t *>(cuSeqlensK_->GetData()) : nullptr;
129	for (int i = 0; i < batchSize; i++) {	147	for (int i = 0; i < batchSize; i++) {
130	if (sequsedKPtr[i] < 0) {	148	if (sequsedKPtr[i] < 0) {
131	KERNEL_LOG_ERROR("The elements in seqused_k should be >= 0, but got seqused_k[%d] = %d", i,	149	KERNEL_LOG_ERROR("The elements in seqused_k should be >= 0, but got seqused_k[%d] = %d", i,
132	sequsedKPtr[i]);	150	sequsedKPtr[i]);
133	return false;	151	return false;
134	}	152	}
		153	+ // 校验 seqused_k 元素不大于 max_seqlen_k (BSND) 或 cu_seqlens_k 序列长度 (TND)
		154	+ if (layoutK_ == "BSND" && sequsedKPtr[i] > maxSeqlenK_) {
		155	+ KERNEL_LOG_ERROR("The elements in seqused_k should not be greater than max_seqlen_k %d, "
		156	+ "but got seqused_k[%d] = %d", maxSeqlenK_, i, sequsedKPtr[i]);
		157	+ return false;
		158	+ }
		159	+ if (cuSeqlensKPtr != nullptr) {
		160	+ int32_t seqLen = cuSeqlensKPtr[i + 1] - cuSeqlensKPtr[i];
		161	+ if (sequsedKPtr[i] > seqLen) {
		162	+ KERNEL_LOG_ERROR("The elements in seqused_k should not be greater than the sequence length "
		163	+ "from cu_seqlens_k %d, but got seqused_k[%d] = %d", seqLen, i, sequsedKPtr[i]);
		164	+ return false;
		165	+ }
		166	+ }
135	}	167	}
136	}	168	}
137	- // 校验 cmp_residual_k 元素非负	169	+ // 校验 cmp_residual_k 元素
138	if (cmpResidualK_ != nullptr && cmpResidualK_->GetData() != nullptr) {	170	if (cmpResidualK_ != nullptr && cmpResidualK_->GetData() != nullptr) {
139	const int32_t cmpResidualKPtr = static_cast<const int32_t >(cmpResidualK_->GetData());	171	const int32_t cmpResidualKPtr = static_cast<const int32_t >(cmpResidualK_->GetData());
140	for (int i = 0; i < batchSize; i++) {	172	for (int i = 0; i < batchSize; i++) {
141	if (cmpResidualKPtr[i] < 0 \|\| cmpResidualKPtr[i] >= cmpRatio_) {	173	if (cmpResidualKPtr[i] < 0 \|\| cmpResidualKPtr[i] >= cmpRatio_) {
142	- KERNEL_LOG_ERROR("The elements in cmp_residual_k should be in [0, cmpRatio_), but got "	174	+ KERNEL_LOG_ERROR("The elements in cmp_residual_k should be in [0, cmpRatio_(%d)), but got "
143	- "cmp_residual_k[%d] = %d",	175	+ "cmp_residual_k[%d] = %d", cmpRatio_,
144	i, cmpResidualKPtr[i]);	176	i, cmpResidualKPtr[i]);
145	return false;	177	return false;
146	}	178	}

Mattention/sparse_flash_mla_grad_metadata/op_host/sparse_flash_mla_grad_metadata_check.h+160-3

	@@ -222,6 +222,31 @@ aclnnStatus CheckSingleParamSmlag(int64_t batchSize, int64_t maxSeqlenQ, int64_t
222	"The value of layout_q must be equal to that of layout_kv");	222	"The value of layout_q must be equal to that of layout_kv");
223	return ACLNN_ERR_PARAM_INVALID;	223	return ACLNN_ERR_PARAM_INVALID;
224	}	224	}
		225	+ // 校验 layout_q 为 BSND 时，max_seqlen_q 必须大于 0
		226	+ if (strcmp(layoutQOptional, "BSND") == 0 && maxSeqlenQ <= 0) {
		227	+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(SMLAG_ACLNN_OP_NAME, "max_seqlen_q", std::to_string(maxSeqlenQ),
		228	+ "When layout_q is BSND, the value of max_seqlen_q "
		229	+ "must be equal to the size of the second axis of q");
		230	+ return ACLNN_ERR_PARAM_INVALID;
		231	+ }
		232	+ // 校验 has_ori_kv 且 layout_kv 为 BSND 时，max_seqlen_ori_kv 必须大于 0
		233	+ if (hasOriKv && strcmp(layoutKvOptional, "BSND") == 0 && maxSeqlenOriKv <= 0) {
		234	+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(SMLAG_ACLNN_OP_NAME, "max_seqlen_ori_kv",
		235	+ std::to_string(maxSeqlenOriKv),
		236	+ "When has_ori_kv is true and layout_kv is BSND, "
		237	+ "the value of max_seqlen_ori_kv "
		238	+ "must be equal to the size of the second axis of ori_kv");
		239	+ return ACLNN_ERR_PARAM_INVALID;
		240	+ }
		241	+ // 校验 has_cmp_kv 且 layout_kv 为 BSND 时，max_seqlen_cmp_kv 必须大于 0
		242	+ if (hasCmpKv && strcmp(layoutKvOptional, "BSND") == 0 && maxSeqlenCmpKv <= 0) {
		243	+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(SMLAG_ACLNN_OP_NAME, "max_seqlen_cmp_kv",
		244	+ std::to_string(maxSeqlenCmpKv),
		245	+ "When has_cmp_kv is true and layout_kv is BSND, "
		246	+ "the value of max_seqlen_cmp_kv "
		247	+ "must be equal to the size of the second axis of cmp_kv");
		248	+ return ACLNN_ERR_PARAM_INVALID;
		249	+ }
225	// 核数校验	250	// 核数校验
226	if (aicCoreNum == 0) {	251	if (aicCoreNum == 0) {
227	OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(SMLAG_ACLNN_OP_NAME, "aic_core_num", std::to_string(aicCoreNum),	252	OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(SMLAG_ACLNN_OP_NAME, "aic_core_num", std::to_string(aicCoreNum),
	@@ -360,13 +385,72 @@ int64_t GetCmpKvBatchSizeSmlag(const aclTensor *sequsedCmpKvOptional, const aclT
360	return batchSize;	385	return batchSize;
361	}	386	}
362		387
		388	+std::string TopkLengthShapeToStringSmlag(const aclTensor *topkLengthOptional)
		389	+{
		390	+ const auto &shape = topkLengthOptional->GetViewShape();
		391	+ std::string result;
		392	+ for (size_t i = 0; i < shape.GetDimNum(); ++i) {
		393	+ if (i != 0) {
		394	+ result += ", ";
		395	+ }
		396	+ result += std::to_string(shape.GetDim(i));
		397	+ }
		398	+ return result;
		399	+}
		400	+
		401	+aclnnStatus CheckTopkLengthFirstDimSmlag(const aclTensor *topkLengthOptional, const std::string &topkLengthName,
		402	+ int64_t queryBatchSize, const std::string &querySource)
		403	+{
		404	+ if (topkLengthOptional->GetViewShape().GetDim(0) == queryBatchSize) {
		405	+ return ACLNN_SUCCESS;
		406	+ }
		407	+ std::string incorrectShape = TopkLengthShapeToStringSmlag(topkLengthOptional);
		408	+ if (IsTensorSourceSmlag(querySource)) {
		409	+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
		410	+ SMLAG_ACLNN_OP_NAME, topkLengthName, incorrectShape,
		411	+ "When layout_q is BSND, the size of the first axis of " + topkLengthName +
		412	+ " must be equal to " + GetSourceDescSmlag(querySource));
		413	+ } else {
		414	+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
		415	+ SMLAG_ACLNN_OP_NAME, topkLengthName, incorrectShape,
		416	+ "When layout_q is BSND, the size of the first axis of " + topkLengthName +
		417	+ " must be equal to batch_size");
		418	+ }
		419	+ return ACLNN_ERR_PARAM_INVALID;
		420	+}
		421	+
		422	+struct TopkLengthAxisSmlag {
		423	+ int64_t index;
		424	+ const char *desc;
		425	+};
		426	+
		427	+inline constexpr TopkLengthAxisSmlag SMLAG_TOPK_LENGTH_SECOND_AXIS{1, "second"};
		428	+inline constexpr TopkLengthAxisSmlag SMLAG_TOPK_LENGTH_THIRD_AXIS{2, "third"};
		429	+
		430	+aclnnStatus CheckTopkLengthSingleDimSmlag(const aclTensor *topkLengthOptional, const std::string &topkLengthName,
		431	+ TopkLengthAxisSmlag axis, int64_t expectedValue,
		432	+ const std::string &expectedDesc, const char *layoutQOptional)
		433	+{
		434	+ if (topkLengthOptional->GetViewShape().GetDim(axis.index) == expectedValue) {
		435	+ return ACLNN_SUCCESS;
		436	+ }
		437	+ std::string incorrectShape = TopkLengthShapeToStringSmlag(topkLengthOptional);
		438	+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
		439	+ SMLAG_ACLNN_OP_NAME, topkLengthName, incorrectShape,
		440	+ "When layout_q is " + std::string(layoutQOptional) + ", the size of the " + axis.desc + " axis of " +
		441	+ topkLengthName + " must be equal to " + expectedDesc);
		442	+ return ACLNN_ERR_PARAM_INVALID;
		443	+}
		444	+
363	aclnnStatus CheckConsistencySmlag(const aclTensor cuSeqlensQOptional, const aclTensor cuSeqlensOriKvOptional,	445	aclnnStatus CheckConsistencySmlag(const aclTensor cuSeqlensQOptional, const aclTensor cuSeqlensOriKvOptional,
364	const aclTensor cuSeqlensCmpKvOptional, const aclTensor sequsedQOptional,	446	const aclTensor cuSeqlensCmpKvOptional, const aclTensor sequsedQOptional,
365	const aclTensor sequsedOriKvOptional, const aclTensor sequsedCmpKvOptional,	447	const aclTensor sequsedOriKvOptional, const aclTensor sequsedCmpKvOptional,
366	const aclTensor cmpResidualKvOptional, const aclTensor oriTopkLengthOptional,	448	const aclTensor cmpResidualKvOptional, const aclTensor oriTopkLengthOptional,
367	const aclTensor *cmpTopkLengthOptional, int64_t batchSize,	449	const aclTensor *cmpTopkLengthOptional, int64_t batchSize,
368	const char layoutQOptional, const char layoutKvOptional, bool hasOriKv,	450	const char layoutQOptional, const char layoutKvOptional, bool hasOriKv,
369	- bool hasCmpKv, const aclTensor *metadata)	451	+ bool hasCmpKv, int64_t oriTopk, int64_t cmpTopk, int64_t oriMaskMode,
		452	+ int64_t cmpMaskMode, int64_t maxSeqlenQ, int64_t numHeadsKv,
		453	+ const aclTensor *metadata)
370	{	454	{
371	aclDataType dataType = aclDataType::ACL_DT_UNDEFINED;	455	aclDataType dataType = aclDataType::ACL_DT_UNDEFINED;
372	int64_t dimNum = -1;	456	int64_t dimNum = -1;
	@@ -441,7 +525,8 @@ aclnnStatus CheckConsistencySmlag(const aclTensor *cuSeqlensQOptional, const acl
441	}	525	}
442	}	526	}
443	// 校验 ori_topk_length	527	// 校验 ori_topk_length
444	- if (IsTensorExistSmlag(oriTopkLengthOptional)) {	528	+ if (oriTopk != 0 && oriMaskMode == static_cast<int64_t>(SparseModeSmlag::DEFAULT_MASK) &&
		529	+ IsTensorExistSmlag(oriTopkLengthOptional)) {
445	// 校验 ori_topk_length 维度	530	// 校验 ori_topk_length 维度
446	dimNum = GetDimNumSmlag(oriTopkLengthOptional);	531	dimNum = GetDimNumSmlag(oriTopkLengthOptional);
447	if (strcmp(layoutQOptional, "TND") == 0) {	532	if (strcmp(layoutQOptional, "TND") == 0) {
	@@ -528,7 +613,8 @@ aclnnStatus CheckConsistencySmlag(const aclTensor *cuSeqlensQOptional, const acl
528	}	613	}
529	}	614	}
530	// 校验 cmp_topk_length	615	// 校验 cmp_topk_length
531	- if (IsTensorExistSmlag(cmpTopkLengthOptional)) {	616	+ if (cmpTopk != 0 && cmpMaskMode == static_cast<int64_t>(SparseModeSmlag::DEFAULT_MASK) &&
		617	+ IsTensorExistSmlag(cmpTopkLengthOptional)) {
532	// 校验 cmp_topk_length 维度	618	// 校验 cmp_topk_length 维度
533	dimNum = GetDimNumSmlag(cmpTopkLengthOptional);	619	dimNum = GetDimNumSmlag(cmpTopkLengthOptional);
534	if (strcmp(layoutQOptional, "TND") == 0) {	620	if (strcmp(layoutQOptional, "TND") == 0) {
	@@ -639,6 +725,41 @@ aclnnStatus CheckConsistencySmlag(const aclTensor *cuSeqlensQOptional, const acl
639	return ACLNN_ERR_PARAM_INVALID;	725	return ACLNN_ERR_PARAM_INVALID;
640	}	726	}
641	}	727	}
		728	+ // 校验 ori_topk_length 维度一致性
		729	+ if (oriTopk != 0 &&
		730	+ oriMaskMode == static_cast<int64_t>(SparseModeSmlag::DEFAULT_MASK) &&
		731	+ IsTensorExistSmlag(oriTopkLengthOptional)) {
		732	+ if (strcmp(layoutQOptional, "BSND") == 0) {
		733	+ // 校验 ori_topk_length 第一个维度
		734	+ aclnnStatus ret = CheckTopkLengthFirstDimSmlag(oriTopkLengthOptional, "ori_topk_length",
		735	+ queryBatchSize, querySource);
		736	+ if (ret != ACLNN_SUCCESS) {
		737	+ return ret;
		738	+ }
		739	+ // 校验 ori_topk_length 第二个维度
		740	+ ret = CheckTopkLengthSingleDimSmlag(oriTopkLengthOptional, "ori_topk_length",
		741	+ SMLAG_TOPK_LENGTH_SECOND_AXIS, maxSeqlenQ, "max_seqlen_q",
		742	+ layoutQOptional);
		743	+ if (ret != ACLNN_SUCCESS) {
		744	+ return ret;
		745	+ }
		746	+ // 校验 ori_topk_length 第三个维度
		747	+ ret = CheckTopkLengthSingleDimSmlag(oriTopkLengthOptional, "ori_topk_length",
		748	+ SMLAG_TOPK_LENGTH_THIRD_AXIS, numHeadsKv, "num_heads_kv",
		749	+ layoutQOptional);
		750	+ if (ret != ACLNN_SUCCESS) {
		751	+ return ret;
		752	+ }
		753	+ } else if (strcmp(layoutQOptional, "TND") == 0) {
		754	+ // 校验 ori_topk_length 第二个维度
		755	+ aclnnStatus ret = CheckTopkLengthSingleDimSmlag(oriTopkLengthOptional, "ori_topk_length",
		756	+ SMLAG_TOPK_LENGTH_SECOND_AXIS, numHeadsKv,
		757	+ "num_heads_kv", layoutQOptional);
		758	+ if (ret != ACLNN_SUCCESS) {
		759	+ return ret;
		760	+ }
		761	+ }
		762	+ }
642	}	763	}
643	if (hasCmpKv) {	764	if (hasCmpKv) {
644	std::string cmpKvSource;	765	std::string cmpKvSource;
	@@ -699,6 +820,41 @@ aclnnStatus CheckConsistencySmlag(const aclTensor *cuSeqlensQOptional, const acl
699	return ACLNN_ERR_PARAM_INVALID;	820	return ACLNN_ERR_PARAM_INVALID;
700	}	821	}
701	}	822	}
		823	+ // 校验 cmp_topk_length 维度一致性
		824	+ if (cmpTopk != 0 &&
		825	+ cmpMaskMode == static_cast<int64_t>(SparseModeSmlag::DEFAULT_MASK) &&
		826	+ IsTensorExistSmlag(cmpTopkLengthOptional)) {
		827	+ if (strcmp(layoutQOptional, "BSND") == 0) {
		828	+ // 校验 cmp_topk_length 第一个维度
		829	+ aclnnStatus ret = CheckTopkLengthFirstDimSmlag(cmpTopkLengthOptional, "cmp_topk_length",
		830	+ queryBatchSize, querySource);
		831	+ if (ret != ACLNN_SUCCESS) {
		832	+ return ret;
		833	+ }
		834	+ // 校验 cmp_topk_length 第二个维度
		835	+ ret = CheckTopkLengthSingleDimSmlag(cmpTopkLengthOptional, "cmp_topk_length",
		836	+ SMLAG_TOPK_LENGTH_SECOND_AXIS, maxSeqlenQ, "max_seqlen_q",
		837	+ layoutQOptional);
		838	+ if (ret != ACLNN_SUCCESS) {
		839	+ return ret;
		840	+ }
		841	+ // 校验 cmp_topk_length 第三个维度
		842	+ ret = CheckTopkLengthSingleDimSmlag(cmpTopkLengthOptional, "cmp_topk_length",
		843	+ SMLAG_TOPK_LENGTH_THIRD_AXIS, numHeadsKv, "num_heads_kv",
		844	+ layoutQOptional);
		845	+ if (ret != ACLNN_SUCCESS) {
		846	+ return ret;
		847	+ }
		848	+ } else if (strcmp(layoutQOptional, "TND") == 0) {
		849	+ // 校验 cmp_topk_length 第二个维度
		850	+ aclnnStatus ret = CheckTopkLengthSingleDimSmlag(cmpTopkLengthOptional, "cmp_topk_length",
		851	+ SMLAG_TOPK_LENGTH_SECOND_AXIS, numHeadsKv,
		852	+ "num_heads_kv", layoutQOptional);
		853	+ if (ret != ACLNN_SUCCESS) {
		854	+ return ret;
		855	+ }
		856	+ }
		857	+ }
702	}	858	}
703	return ACLNN_SUCCESS;	859	return ACLNN_SUCCESS;
704	}	860	}
	@@ -726,6 +882,7 @@ static aclnnStatus ParamsCheck(const aclTensor *cuSeqlensQOptional, const aclTen
726	CheckConsistencySmlag(cuSeqlensQOptional, cuSeqlensOriKvOptional, cuSeqlensCmpKvOptional, sequsedQOptional,	882	CheckConsistencySmlag(cuSeqlensQOptional, cuSeqlensOriKvOptional, cuSeqlensCmpKvOptional, sequsedQOptional,
727	sequsedOriKvOptional, sequsedCmpKvOptional, cmpResidualKvOptional, oriTopkLengthOptional,	883	sequsedOriKvOptional, sequsedCmpKvOptional, cmpResidualKvOptional, oriTopkLengthOptional,
728	cmpTopkLengthOptional, batchSize, layoutQOptional, layoutKvOptional, hasOriKv, hasCmpKv,	884	cmpTopkLengthOptional, batchSize, layoutQOptional, layoutKvOptional, hasOriKv, hasCmpKv,
		885	+ oriTopk, cmpTopk, oriMaskMode, cmpMaskMode, maxSeqlenQ, numHeadsKv,
729	metaData) == ACLNN_SUCCESS) {	886	metaData) == ACLNN_SUCCESS) {
730	return ACLNN_SUCCESS;	887	return ACLNN_SUCCESS;
731	} else {	888	} else {

Mattention/sparse_flash_mla_grad_metadata/op_kernel_aicpu/sparse_flash_mla_grad_metadata_aicpu.cpp+85-35

	@@ -87,13 +87,12 @@ bool SparseFlashMlaGradMetadataCpuKernel::ParamsCheck()
87	if (layoutQ_ == "TND") {	87	if (layoutQ_ == "TND") {
88	if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) {	88	if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) {
89	const int32_t cuSeqlensQPtr = static_cast<const int32_t >(cuSeqlensQ_->GetData());	89	const int32_t cuSeqlensQPtr = static_cast<const int32_t >(cuSeqlensQ_->GetData());
		90	+ // 校验 cu_seqlens_q 首元素为 0
		91	+ if (cuSeqlensQPtr[0] != 0) {
		92	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_q should be 0, but got %d", cuSeqlensQPtr[0]);
		93	+ return false;
		94	+ }
90	for (int i = 0; i < batchSize + 1; i++) {	95	for (int i = 0; i < batchSize + 1; i++) {
91	- // 校验 cu_seqlens_q 元素非负
92	- if (cuSeqlensQPtr[i] < 0) {
93	- KERNEL_LOG_ERROR("The elements in cu_seqlens_q should be >= 0, but got cu_seqlens_q[%d] = %d", i,
94	- cuSeqlensQPtr[i]);
95	- return false;
96	- }
97	// 校验 cu_seqlens_q 元素递增	96	// 校验 cu_seqlens_q 元素递增
98	if (i > 0 && cuSeqlensQPtr[i - 1] > cuSeqlensQPtr[i]) {	97	if (i > 0 && cuSeqlensQPtr[i - 1] > cuSeqlensQPtr[i]) {
99	KERNEL_LOG_ERROR("The elements in cu_seqlens_q must be in ascending order, "	98	KERNEL_LOG_ERROR("The elements in cu_seqlens_q must be in ascending order, "
	@@ -107,6 +106,9 @@ bool SparseFlashMlaGradMetadataCpuKernel::ParamsCheck()
107	// 校验 seqused_q 元素	106	// 校验 seqused_q 元素
108	if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) {	107	if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) {
109	const int32_t sequsedQPtr = static_cast<const int32_t >(sequsedQ_->GetData());	108	const int32_t sequsedQPtr = static_cast<const int32_t >(sequsedQ_->GetData());
		109	+ const int32_t *cuSeqlensQPtr = (layoutQ_ == "TND" && cuSeqlensQ_ != nullptr &&
		110	+ cuSeqlensQ_->GetData() != nullptr) ?
		111	+ static_cast<const int32_t *>(cuSeqlensQ_->GetData()) : nullptr;
110	for (int i = 0; i < batchSize; i++) {	112	for (int i = 0; i < batchSize; i++) {
111	// 校验 seqused_q 元素非负	113	// 校验 seqused_q 元素非负
112	if (sequsedQPtr[i] < 0) {	114	if (sequsedQPtr[i] < 0) {
	@@ -114,6 +116,20 @@ bool SparseFlashMlaGradMetadataCpuKernel::ParamsCheck()
114	sequsedQPtr[i]);	116	sequsedQPtr[i]);
115	return false;	117	return false;
116	}	118	}
		119	+ // 校验 seqused_q 元素不大于 max_seqlen_q (BSND) 或 cu_seqlens_q 序列长度 (TND)
		120	+ if (layoutQ_ == "BSND" && sequsedQPtr[i] > maxSeqlenQ_) {
		121	+ KERNEL_LOG_ERROR("The elements in seqused_q should not be greater than max_seqlen_q %d, "
		122	+ "but got seqused_q[%d] = %d", maxSeqlenQ_, i, sequsedQPtr[i]);
		123	+ return false;
		124	+ }
		125	+ if (cuSeqlensQPtr != nullptr) {
		126	+ int32_t seqLen = cuSeqlensQPtr[i + 1] - cuSeqlensQPtr[i];
		127	+ if (sequsedQPtr[i] > seqLen) {
		128	+ KERNEL_LOG_ERROR("The elements in seqused_q should not be greater than the sequence length "
		129	+ "from cu_seqlens_q %d, but got seqused_q[%d] = %d", seqLen, i, sequsedQPtr[i]);
		130	+ return false;
		131	+ }
		132	+ }
117	}	133	}
118	}	134	}
119	if (hasOriKv_) {	135	if (hasOriKv_) {
	@@ -121,14 +137,13 @@ bool SparseFlashMlaGradMetadataCpuKernel::ParamsCheck()
121	if (layoutKv_ == "TND") {	137	if (layoutKv_ == "TND") {
122	if (cuSeqlensOriKv_ != nullptr && cuSeqlensOriKv_->GetData() != nullptr) {	138	if (cuSeqlensOriKv_ != nullptr && cuSeqlensOriKv_->GetData() != nullptr) {
123	const int32_t cuSeqlensOriKvPtr = static_cast<const int32_t >(cuSeqlensOriKv_->GetData());	139	const int32_t cuSeqlensOriKvPtr = static_cast<const int32_t >(cuSeqlensOriKv_->GetData());
		140	+ // 校验 cu_seqlens_ori_kv 首元素为 0
		141	+ if (cuSeqlensOriKvPtr[0] != 0) {
		142	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_ori_kv should be 0, but got %d",
		143	+ cuSeqlensOriKvPtr[0]);
		144	+ return false;
		145	+ }
124	for (int i = 0; i < batchSize + 1; i++) {	146	for (int i = 0; i < batchSize + 1; i++) {
125	- // 校验 cu_seqlens_ori_kv 元素非负
126	- if (cuSeqlensOriKvPtr[i] < 0) {
127	- KERNEL_LOG_ERROR("The elements in cu_seqlens_ori_kv should be >= 0, "
128	- "but got cu_seqlens_ori_kv[%d] = %d",
129	- i, cuSeqlensOriKvPtr[i]);
130	- return false;
131	- }
132	// 校验 cu_seqlens_ori_kv 元素递增	147	// 校验 cu_seqlens_ori_kv 元素递增
133	if (i > 0 && cuSeqlensOriKvPtr[i - 1] > cuSeqlensOriKvPtr[i]) {	148	if (i > 0 && cuSeqlensOriKvPtr[i - 1] > cuSeqlensOriKvPtr[i]) {
134	KERNEL_LOG_ERROR("The elements in cu_seqlens_ori_kv must be in ascending order, "	149	KERNEL_LOG_ERROR("The elements in cu_seqlens_ori_kv must be in ascending order, "
	@@ -142,6 +157,10 @@ bool SparseFlashMlaGradMetadataCpuKernel::ParamsCheck()
142	// 校验 seqused_ori_kv 元素	157	// 校验 seqused_ori_kv 元素
143	if (sequsedOriKv_ != nullptr && sequsedOriKv_->GetData() != nullptr) {	158	if (sequsedOriKv_ != nullptr && sequsedOriKv_->GetData() != nullptr) {
144	const int32_t sequsedOriKvPtr = static_cast<const int32_t >(sequsedOriKv_->GetData());	159	const int32_t sequsedOriKvPtr = static_cast<const int32_t >(sequsedOriKv_->GetData());
		160	+ const int32_t *cuSeqlensOriKvPtr = (layoutKv_ == "TND" && cuSeqlensOriKv_ != nullptr &&
		161	+ cuSeqlensOriKv_->GetData() != nullptr) ?
		162	+ static_cast<const int32_t *>(cuSeqlensOriKv_->GetData()) :
		163	+ nullptr;
145	for (int i = 0; i < batchSize; i++) {	164	for (int i = 0; i < batchSize; i++) {
146	// 校验 seqused_ori_kv 元素非负	165	// 校验 seqused_ori_kv 元素非负
147	if (sequsedOriKvPtr[i] < 0) {	166	if (sequsedOriKvPtr[i] < 0) {
	@@ -149,10 +168,27 @@ bool SparseFlashMlaGradMetadataCpuKernel::ParamsCheck()
149	i, sequsedOriKvPtr[i]);	168	i, sequsedOriKvPtr[i]);
150	return false;	169	return false;
151	}	170	}
		171	+ // 校验 seqused_ori_kv 元素不大于 max_seqlen_ori_kv (BSND) 或 cu_seqlens_ori_kv 序列长度 (TND)
		172	+ if (layoutKv_ == "BSND" && sequsedOriKvPtr[i] > maxSeqlenOriKv_) {
		173	+ KERNEL_LOG_ERROR("The elements in seqused_ori_kv should not be greater than "
		174	+ "max_seqlen_ori_kv %d, but got seqused_ori_kv[%d] = %d",
		175	+ maxSeqlenOriKv_, i, sequsedOriKvPtr[i]);
		176	+ return false;
		177	+ }
		178	+ if (cuSeqlensOriKvPtr != nullptr) {
		179	+ int32_t seqLen = cuSeqlensOriKvPtr[i + 1] - cuSeqlensOriKvPtr[i];
		180	+ if (sequsedOriKvPtr[i] > seqLen) {
		181	+ KERNEL_LOG_ERROR("The elements in seqused_ori_kv should not be greater than the sequence "
		182	+ "length from cu_seqlens_ori_kv %d, but got seqused_ori_kv[%d] = %d",
		183	+ seqLen, i, sequsedOriKvPtr[i]);
		184	+ return false;
		185	+ }
		186	+ }
152	}	187	}
153	}	188	}
154	// 校验 ori_topk_length 元素	189	// 校验 ori_topk_length 元素
155	- if (oriTopkLength_ != nullptr && oriTopkLength_->GetData() != nullptr) {	190	+ if (oriTopK_ != 0 && oriMaskMode_ == static_cast<int32_t>(SparseMode::DEFAULT_MASK) &&
		191	+ oriTopkLength_ != nullptr && oriTopkLength_->GetData() != nullptr) {
156	// 校验 ori_topk_length 元素数量	192	// 校验 ori_topk_length 元素数量
157	int32_t sumOfQuerySeq = GetSumOfQuerySeq();	193	int32_t sumOfQuerySeq = GetSumOfQuerySeq();
158	const int32_t oriTopkLengthPtr = static_cast<const int32_t >(oriTopkLength_->GetData());	194	const int32_t oriTopkLengthPtr = static_cast<const int32_t >(oriTopkLength_->GetData());
	@@ -182,14 +218,13 @@ bool SparseFlashMlaGradMetadataCpuKernel::ParamsCheck()
182	// 校验 cu_seqlens_cmp_kv 元素	218	// 校验 cu_seqlens_cmp_kv 元素
183	if (cuSeqlensCmpKv_ != nullptr && cuSeqlensCmpKv_->GetData() != nullptr) {	219	if (cuSeqlensCmpKv_ != nullptr && cuSeqlensCmpKv_->GetData() != nullptr) {
184	const int32_t cuSeqlensCmpKvPtr = static_cast<const int32_t >(cuSeqlensCmpKv_->GetData());	220	const int32_t cuSeqlensCmpKvPtr = static_cast<const int32_t >(cuSeqlensCmpKv_->GetData());
		221	+ // 校验 cu_seqlens_cmp_kv 首元素为 0
		222	+ if (cuSeqlensCmpKvPtr[0] != 0) {
		223	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_cmp_kv should be 0, but got %d",
		224	+ cuSeqlensCmpKvPtr[0]);
		225	+ return false;
		226	+ }
185	for (int i = 0; i < batchSize + 1; i++) {	227	for (int i = 0; i < batchSize + 1; i++) {
186	- // 校验 cu_seqlens_cmp_kv 元素非负
187	- if (cuSeqlensCmpKvPtr[i] < 0) {
188	- KERNEL_LOG_ERROR("The elements in cu_seqlens_cmp_kv should be >= 0, "
189	- "but got cu_seqlens_cmp_kv[%d] = %d",
190	- i, cuSeqlensCmpKvPtr[i]);
191	- return false;
192	- }
193	// 校验 cu_seqlens_cmp_kv 元素递增	228	// 校验 cu_seqlens_cmp_kv 元素递增
194	if (i > 0 && cuSeqlensCmpKvPtr[i - 1] > cuSeqlensCmpKvPtr[i]) {	229	if (i > 0 && cuSeqlensCmpKvPtr[i - 1] > cuSeqlensCmpKvPtr[i]) {
195	KERNEL_LOG_ERROR("The elements in cu_seqlens_cmp_kv must be in ascending order, "	230	KERNEL_LOG_ERROR("The elements in cu_seqlens_cmp_kv must be in ascending order, "
	@@ -203,6 +238,10 @@ bool SparseFlashMlaGradMetadataCpuKernel::ParamsCheck()
203	// 校验 seqused_cmp_kv 元素	238	// 校验 seqused_cmp_kv 元素
204	if (sequsedCmpKv_ != nullptr && sequsedCmpKv_->GetData() != nullptr) {	239	if (sequsedCmpKv_ != nullptr && sequsedCmpKv_->GetData() != nullptr) {
205	const int32_t sequsedCmpKvPtr = static_cast<const int32_t >(sequsedCmpKv_->GetData());	240	const int32_t sequsedCmpKvPtr = static_cast<const int32_t >(sequsedCmpKv_->GetData());
		241	+ const int32_t *cuSeqlensCmpKvPtr = (layoutKv_ == "TND" && cuSeqlensCmpKv_ != nullptr &&
		242	+ cuSeqlensCmpKv_->GetData() != nullptr) ?
		243	+ static_cast<const int32_t *>(cuSeqlensCmpKv_->GetData()) :
		244	+ nullptr;
206	for (int i = 0; i < batchSize; i++) {	245	for (int i = 0; i < batchSize; i++) {
207	// 校验 seqused_cmp_kv 元素非负	246	// 校验 seqused_cmp_kv 元素非负
208	if (sequsedCmpKvPtr[i] < 0) {	247	if (sequsedCmpKvPtr[i] < 0) {
	@@ -210,22 +249,39 @@ bool SparseFlashMlaGradMetadataCpuKernel::ParamsCheck()
210	i, sequsedCmpKvPtr[i]);	249	i, sequsedCmpKvPtr[i]);
211	return false;	250	return false;
212	}	251	}
		252	+ // 校验 seqused_cmp_kv 元素不大于 max_seqlen_cmp_kv (BSND) 或 cu_seqlens_cmp_kv 序列长度 (TND)
		253	+ if (layoutKv_ == "BSND" && sequsedCmpKvPtr[i] > maxSeqlenCmpKv_) {
		254	+ KERNEL_LOG_ERROR("The elements in seqused_cmp_kv should not be greater than "
		255	+ "max_seqlen_cmp_kv %d, but got seqused_cmp_kv[%d] = %d",
		256	+ maxSeqlenCmpKv_, i, sequsedCmpKvPtr[i]);
		257	+ return false;
		258	+ }
		259	+ if (cuSeqlensCmpKvPtr != nullptr) {
		260	+ int32_t seqLen = cuSeqlensCmpKvPtr[i + 1] - cuSeqlensCmpKvPtr[i];
		261	+ if (sequsedCmpKvPtr[i] > seqLen) {
		262	+ KERNEL_LOG_ERROR("The elements in seqused_cmp_kv should not be greater than the sequence "
		263	+ "length from cu_seqlens_cmp_kv %d, but got seqused_cmp_kv[%d] = %d",
		264	+ seqLen, i, sequsedCmpKvPtr[i]);
		265	+ return false;
		266	+ }
		267	+ }
213	}	268	}
214	}	269	}
215	// 校验 cmp_residual_kv 元素	270	// 校验 cmp_residual_kv 元素
216	if (cmpResidualKv_ != nullptr && cmpResidualKv_->GetData() != nullptr) {	271	if (cmpResidualKv_ != nullptr && cmpResidualKv_->GetData() != nullptr) {
217	const int32_t cmpResidualKvPtr = static_cast<const int32_t >(cmpResidualKv_->GetData());	272	const int32_t cmpResidualKvPtr = static_cast<const int32_t >(cmpResidualKv_->GetData());
218	for (int i = 0; i < batchSize; i++) {	273	for (int i = 0; i < batchSize; i++) {
219	- // 校验 cmp_residual_kv 元素非负	274	+ if (cmpResidualKvPtr[i] < 0 \|\| cmpResidualKvPtr[i] >= cmpRatio_) {
220	- if (cmpResidualKvPtr[i] < 0) {	275	+ KERNEL_LOG_ERROR("The elements in cmp_residual_kv should be in [0, cmpRatio_(%d)), but got "
221	- KERNEL_LOG_ERROR("The elements in cmp_residual_kv should be >= 0, but got cmp_residual_kv[%d] = %d",	276	+ "cmp_residual_kv[%d] = %d", cmpRatio_,
222	i, cmpResidualKvPtr[i]);	277	i, cmpResidualKvPtr[i]);
223	return false;	278	return false;
224	}	279	}
225	}	280	}
226	}	281	}
227	// 校验 cmp_topk_length 元素	282	// 校验 cmp_topk_length 元素
228	- if (cmpTopkLength_ != nullptr && cmpTopkLength_->GetData() != nullptr) {	283	+ if (cmpTopK_ != 0 && cmpMaskMode_ == static_cast<int32_t>(SparseMode::DEFAULT_MASK) &&
		284	+ cmpTopkLength_ != nullptr && cmpTopkLength_->GetData() != nullptr) {
229	// 校验 cmp_topk_length 元素数量	285	// 校验 cmp_topk_length 元素数量
230	int32_t sumOfQuerySeq = GetSumOfQuerySeq();	286	int32_t sumOfQuerySeq = GetSumOfQuerySeq();
231	const int32_t cmpTopkLengthPtr = static_cast<const int32_t >(cmpTopkLength_->GetData());	287	const int32_t cmpTopkLengthPtr = static_cast<const int32_t >(cmpTopkLength_->GetData());
	@@ -402,7 +458,8 @@ uint32_t SparseFlashMlaGradMetadataCpuKernel::GetBsStride(uint32_t bIdx, uint32_
402	uint32_t SparseFlashMlaGradMetadataCpuKernel::GetOriTopkLength(uint32_t bsStride)	458	uint32_t SparseFlashMlaGradMetadataCpuKernel::GetOriTopkLength(uint32_t bsStride)
403	{	459	{
404	// 尝试使用 oriTopkLength_	460	// 尝试使用 oriTopkLength_
405	- if (oriTopkLength_ != nullptr && oriTopkLength_->GetData() != nullptr) {	461	+ if (oriTopK_ != 0 && oriMaskMode_ == static_cast<int32_t>(SparseMode::DEFAULT_MASK) &&
		462	+ oriTopkLength_ != nullptr && oriTopkLength_->GetData() != nullptr) {
406	const int32_t oriTopkPtr = static_cast<const int32_t >(oriTopkLength_->GetData());	463	const int32_t oriTopkPtr = static_cast<const int32_t >(oriTopkLength_->GetData());
407	return static_cast<uint32_t>(oriTopkPtr[bsStride]);	464	return static_cast<uint32_t>(oriTopkPtr[bsStride]);
408	}	465	}
	@@ -413,7 +470,8 @@ uint32_t SparseFlashMlaGradMetadataCpuKernel::GetOriTopkLength(uint32_t bsStride
413	uint32_t SparseFlashMlaGradMetadataCpuKernel::GetCmpTopkLength(uint32_t bsStride)	470	uint32_t SparseFlashMlaGradMetadataCpuKernel::GetCmpTopkLength(uint32_t bsStride)
414	{	471	{
415	// 尝试使用 cmpTopkLength_	472	// 尝试使用 cmpTopkLength_
416	- if (cmpTopkLength_ != nullptr && cmpTopkLength_->GetData() != nullptr) {	473	+ if (cmpTopK_ != 0 && cmpMaskMode_ == static_cast<int32_t>(SparseMode::DEFAULT_MASK) &&
		474	+ cmpTopkLength_ != nullptr && cmpTopkLength_->GetData() != nullptr) {
417	const int32_t cmpTopkPtr = static_cast<const int32_t >(cmpTopkLength_->GetData());	475	const int32_t cmpTopkPtr = static_cast<const int32_t >(cmpTopkLength_->GetData());
418	return static_cast<uint32_t>(cmpTopkPtr[bsStride]);	476	return static_cast<uint32_t>(cmpTopkPtr[bsStride]);
419	}	477	}
	@@ -455,10 +513,6 @@ uint32_t SparseFlashMlaGradMetadataCpuKernel::GetOriS2SeqSize(uint32_t bIdx)
455	return static_cast<uint32_t>(s2Ptr[bIdx + 1U] - s2Ptr[bIdx]);	513	return static_cast<uint32_t>(s2Ptr[bIdx + 1U] - s2Ptr[bIdx]);
456	}	514	}
457	}	515	}
458	- // 如果 max_seqlen_ori_kv 没传入，且 ori_kv 为稀疏的，则尝试从 topk 中获取
459	- if (maxSeqlenOriKv_ == 0 && isSparseOriKv_) {
460	- return UINT32_MAX;
461	- }
462	// 使用 max_seqlen_ori_kv	516	// 使用 max_seqlen_ori_kv
463	return static_cast<uint32_t>(maxSeqlenOriKv_);	517	return static_cast<uint32_t>(maxSeqlenOriKv_);
464	}	518	}
	@@ -478,10 +532,6 @@ uint32_t SparseFlashMlaGradMetadataCpuKernel::GetCmpS2SeqSize(uint32_t bIdx)
478	return static_cast<uint32_t>(s2Ptr[bIdx + 1U] - s2Ptr[bIdx]);	532	return static_cast<uint32_t>(s2Ptr[bIdx + 1U] - s2Ptr[bIdx]);
479	}	533	}
480	}	534	}
481	- // 如果 max_seqlen_cmp_kv 没传入，且 cmp_kv 为稀疏的，则尝试从topk中获取
482	- if (maxSeqlenCmpKv_ == 0 && isSparseCmpKv_) {
483	- return UINT32_MAX;
484	- }
485	// 使用 max_seqlen_cmp_kv	535	// 使用 max_seqlen_cmp_kv
486	return static_cast<uint32_t>(maxSeqlenCmpKv_);	536	return static_cast<uint32_t>(maxSeqlenCmpKv_);
487	}	537	}

Mattention/sparse_flash_mla_metadata/op_host/sparse_flash_mla_metadata_check.h+165-9

	@@ -220,8 +220,9 @@ aclnnStatus CheckSingleParamSmla(int64_t batchSize, int64_t maxSeqlenQ, int64_t
220	}	220	}
221	// A5 treats -1 as unlimited window	221	// A5 treats -1 as unlimited window
222	if (oriWinLeft < -1 \|\| oriWinRight < -1) {	222	if (oriWinLeft < -1 \|\| oriWinRight < -1) {
223	- OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(SMLA_ACLNN_OP_NAME, "ori_win_left, ori_win_right",	223	+ OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(SMLA_ACLNN_OP_NAME, "ori_win_left and ori_win_right",
224	- std::to_string(oriWinLeft) + ", " + std::to_string(oriWinRight),	224	+ std::to_string(oriWinLeft) + " and " +
		225	+ std::to_string(oriWinRight),
225	"When has_ori_kv is true, the value of ori_win_left, "	226	"When has_ori_kv is true, the value of ori_win_left, "
226	"ori_win_right must be greater than or equal to -1");	227	"ori_win_right must be greater than or equal to -1");
227	return ACLNN_ERR_PARAM_INVALID;	228	return ACLNN_ERR_PARAM_INVALID;
	@@ -233,8 +234,9 @@ aclnnStatus CheckSingleParamSmla(int64_t batchSize, int64_t maxSeqlenQ, int64_t
233	return ACLNN_ERR_PARAM_INVALID;	234	return ACLNN_ERR_PARAM_INVALID;
234	}	235	}
235	if (oriWinLeft != 127 \|\| oriWinRight != 0) {	236	if (oriWinLeft != 127 \|\| oriWinRight != 0) {
236	- OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(SMLA_ACLNN_OP_NAME, "ori_win_left, ori_win_right",	237	+ OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(SMLA_ACLNN_OP_NAME, "ori_win_left and ori_win_right",
237	- std::to_string(oriWinLeft) + ", " + std::to_string(oriWinRight),	238	+ std::to_string(oriWinLeft) + " and " +
		239	+ std::to_string(oriWinRight),
238	"When has_ori_kv is true, the value of ori_win_left "	240	"When has_ori_kv is true, the value of ori_win_left "
239	"must be 127 and the value of ori_win_right must be 0");	241	"must be 127 and the value of ori_win_right must be 0");
240	return ACLNN_ERR_PARAM_INVALID;	242	return ACLNN_ERR_PARAM_INVALID;
	@@ -331,6 +333,31 @@ aclnnStatus CheckSingleParamSmla(int64_t batchSize, int64_t maxSeqlenQ, int64_t
331	"When layout_kv is not PA_BBND, the values of layout_q, layout_kv must be the same");	333	"When layout_kv is not PA_BBND, the values of layout_q, layout_kv must be the same");
332	return ACLNN_ERR_PARAM_INVALID;	334	return ACLNN_ERR_PARAM_INVALID;
333	}	335	}
		336	+ // 校验 layout_q 为 BSND 时，max_seqlen_q 必须大于 0
		337	+ if (strcmp(layoutQOptional, "BSND") == 0 && maxSeqlenQ <= 0) {
		338	+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(SMLA_ACLNN_OP_NAME, "max_seqlen_q", std::to_string(maxSeqlenQ),
		339	+ "When layout_q is BSND, the value of max_seqlen_q "
		340	+ "must be equal to the size of the second axis of q");
		341	+ return ACLNN_ERR_PARAM_INVALID;
		342	+ }
		343	+ // 校验 has_ori_kv 且 layout_kv 为 BSND 时，max_seqlen_ori_kv 必须大于 0
		344	+ if (hasOriKv && strcmp(layoutKvOptional, "BSND") == 0 && maxSeqlenOriKv <= 0) {
		345	+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(SMLA_ACLNN_OP_NAME, "max_seqlen_ori_kv",
		346	+ std::to_string(maxSeqlenOriKv),
		347	+ "When has_ori_kv is true and layout_kv is BSND, "
		348	+ "the value of max_seqlen_ori_kv "
		349	+ "must be equal to the size of the second axis of ori_kv");
		350	+ return ACLNN_ERR_PARAM_INVALID;
		351	+ }
		352	+ // 校验 has_cmp_kv 且 layout_kv 为 BSND 时，max_seqlen_cmp_kv 必须大于 0
		353	+ if (hasCmpKv && strcmp(layoutKvOptional, "BSND") == 0 && maxSeqlenCmpKv <= 0) {
		354	+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(SMLA_ACLNN_OP_NAME, "max_seqlen_cmp_kv",
		355	+ std::to_string(maxSeqlenCmpKv),
		356	+ "When has_cmp_kv is true and layout_kv is BSND, "
		357	+ "the value of max_seqlen_cmp_kv "
		358	+ "must be equal to the size of the second axis of cmp_kv");
		359	+ return ACLNN_ERR_PARAM_INVALID;
		360	+ }
334	// 核数校验	361	// 核数校验
335	if (aicCoreNum == 0) {	362	if (aicCoreNum == 0) {
336	OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(SMLA_ACLNN_OP_NAME, "aic_core_num", std::to_string(aicCoreNum),	363	OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(SMLA_ACLNN_OP_NAME, "aic_core_num", std::to_string(aicCoreNum),
	@@ -536,13 +563,71 @@ int64_t GetCmpKvBatchSizeSmla(const aclTensor *sequsedCmpKvOptional, const aclTe
536	return batchSize;	563	return batchSize;
537	}	564	}
538		565
		566	+std::string TopkLengthShapeToStringSmla(const aclTensor *topkLengthOptional)
		567	+{
		568	+ const auto &shape = topkLengthOptional->GetViewShape();
		569	+ std::string result;
		570	+ for (size_t i = 0; i < shape.GetDimNum(); ++i) {
		571	+ if (i != 0) {
		572	+ result += ", ";
		573	+ }
		574	+ result += std::to_string(shape.GetDim(i));
		575	+ }
		576	+ return result;
		577	+}
		578	+
		579	+aclnnStatus CheckTopkLengthFirstDimSmla(const aclTensor *topkLengthOptional, const std::string &topkLengthName,
		580	+ int64_t queryBatchSize, const std::string &querySource)
		581	+{
		582	+ if (topkLengthOptional->GetViewShape().GetDim(0) == queryBatchSize) {
		583	+ return ACLNN_SUCCESS;
		584	+ }
		585	+ std::string incorrectShape = TopkLengthShapeToStringSmla(topkLengthOptional);
		586	+ if (IsTensorSourceSmla(querySource)) {
		587	+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
		588	+ SMLA_ACLNN_OP_NAME, topkLengthName, incorrectShape,
		589	+ "When layout_q is BSND, the size of the first axis of " + topkLengthName +
		590	+ " must be equal to " + GetSourceDescSmla(querySource));
		591	+ } else {
		592	+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
		593	+ SMLA_ACLNN_OP_NAME, topkLengthName, incorrectShape,
		594	+ "When layout_q is BSND, the size of the first axis of " + topkLengthName +
		595	+ " must be equal to batch_size");
		596	+ }
		597	+ return ACLNN_ERR_PARAM_INVALID;
		598	+}
		599	+
		600	+struct TopkLengthAxisSmla {
		601	+ int64_t index;
		602	+ const char *desc;
		603	+};
		604	+
		605	+inline constexpr TopkLengthAxisSmla SMLA_TOPK_LENGTH_SECOND_AXIS{1, "second"};
		606	+inline constexpr TopkLengthAxisSmla SMLA_TOPK_LENGTH_THIRD_AXIS{2, "third"};
		607	+
		608	+aclnnStatus CheckTopkLengthSingleDimSmla(const aclTensor *topkLengthOptional, const std::string &topkLengthName,
		609	+ TopkLengthAxisSmla axis, int64_t expectedValue,
		610	+ const std::string &expectedDesc, const char *layoutQOptional)
		611	+{
		612	+ if (topkLengthOptional->GetViewShape().GetDim(axis.index) == expectedValue) {
		613	+ return ACLNN_SUCCESS;
		614	+ }
		615	+ std::string incorrectShape = TopkLengthShapeToStringSmla(topkLengthOptional);
		616	+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
		617	+ SMLA_ACLNN_OP_NAME, topkLengthName, incorrectShape,
		618	+ "When layout_q is " + std::string(layoutQOptional) + ", the size of the " + axis.desc + " axis of " +
		619	+ topkLengthName + " must be equal to " + expectedDesc);
		620	+ return ACLNN_ERR_PARAM_INVALID;
		621	+}
		622	+
539	aclnnStatus CheckConsistencySmla(const aclTensor cuSeqlensQOptional, const aclTensor cuSeqlensOriKvOptional,	623	aclnnStatus CheckConsistencySmla(const aclTensor cuSeqlensQOptional, const aclTensor cuSeqlensOriKvOptional,
540	const aclTensor cuSeqlensCmpKvOptional, const aclTensor sequsedQOptional,	624	const aclTensor cuSeqlensCmpKvOptional, const aclTensor sequsedQOptional,
541	const aclTensor sequsedOriKvOptional, const aclTensor sequsedCmpKvOptional,	625	const aclTensor sequsedOriKvOptional, const aclTensor sequsedCmpKvOptional,
542	const aclTensor cmpResidualKvOptional, const aclTensor oriTopkLengthOptional,	626	const aclTensor cmpResidualKvOptional, const aclTensor oriTopkLengthOptional,
543	const aclTensor cmpTopkLengthOptional, int64_t batchSize, const char layoutQOptional,	627	const aclTensor cmpTopkLengthOptional, int64_t batchSize, const char layoutQOptional,
544	- const char layoutKvOptional, bool hasOriKv, bool hasCmpKv, const char socVersion,	628	+ const char *layoutKvOptional, bool hasOriKv, bool hasCmpKv, int64_t oriTopk,
545	- const aclTensor *metadata)	629	+ int64_t cmpTopk, int64_t oriMaskMode, int64_t cmpMaskMode, int64_t maxSeqlenQ,
		630	+ int64_t numHeadsKv, const char socVersion, const aclTensor metadata)
546	{	631	{
547	aclDataType dataType = aclDataType::ACL_DT_UNDEFINED;	632	aclDataType dataType = aclDataType::ACL_DT_UNDEFINED;
548	int64_t dimNum = -1;	633	int64_t dimNum = -1;
	@@ -622,7 +707,7 @@ aclnnStatus CheckConsistencySmla(const aclTensor *cuSeqlensQOptional, const aclT
622	}	707	}
623	}	708	}
624	// 校验 ori_topk_length	709	// 校验 ori_topk_length
625	- if ((socVersion != nullptr && strstr(socVersion, "Ascend950") != nullptr) &&	710	+ if (oriTopk != 0 && oriMaskMode == static_cast<int64_t>(SparseModeSmla::DEFAULT_MASK) &&
626	IsTensorExistSmla(oriTopkLengthOptional)) {	711	IsTensorExistSmla(oriTopkLengthOptional)) {
627	// 校验 ori_topk_length 维度	712	// 校验 ori_topk_length 维度
628	dimNum = GetDimNumSmla(oriTopkLengthOptional);	713	dimNum = GetDimNumSmla(oriTopkLengthOptional);
	@@ -709,7 +794,7 @@ aclnnStatus CheckConsistencySmla(const aclTensor *cuSeqlensQOptional, const aclT
709	}	794	}
710	}	795	}
711	// 校验 cmp_topk_length	796	// 校验 cmp_topk_length
712	- if ((socVersion != nullptr && strstr(socVersion, "Ascend950") != nullptr) &&	797	+ if (cmpTopk != 0 && cmpMaskMode == static_cast<int64_t>(SparseModeSmla::DEFAULT_MASK) &&
713	IsTensorExistSmla(cmpTopkLengthOptional)) {	798	IsTensorExistSmla(cmpTopkLengthOptional)) {
714	// 校验 cmp_topk_length 维度	799	// 校验 cmp_topk_length 维度
715	dimNum = GetDimNumSmla(cmpTopkLengthOptional);	800	dimNum = GetDimNumSmla(cmpTopkLengthOptional);
	@@ -821,6 +906,41 @@ aclnnStatus CheckConsistencySmla(const aclTensor *cuSeqlensQOptional, const aclT
821	return ACLNN_ERR_PARAM_INVALID;	906	return ACLNN_ERR_PARAM_INVALID;
822	}	907	}
823	}	908	}
		909	+ // 校验 ori_topk_length 维度一致性
		910	+ if (oriTopk != 0 &&
		911	+ oriMaskMode == static_cast<int64_t>(SparseModeSmla::DEFAULT_MASK) &&
		912	+ IsTensorExistSmla(oriTopkLengthOptional)) {
		913	+ if (strcmp(layoutQOptional, "BSND") == 0) {
		914	+ // 校验 ori_topk_length 第一个维度
		915	+ aclnnStatus ret = CheckTopkLengthFirstDimSmla(oriTopkLengthOptional, "ori_topk_length",
		916	+ queryBatchSize, querySource);
		917	+ if (ret != ACLNN_SUCCESS) {
		918	+ return ret;
		919	+ }
		920	+ // 校验 ori_topk_length 第二个维度
		921	+ ret = CheckTopkLengthSingleDimSmla(oriTopkLengthOptional, "ori_topk_length",
		922	+ SMLA_TOPK_LENGTH_SECOND_AXIS, maxSeqlenQ, "max_seqlen_q",
		923	+ layoutQOptional);
		924	+ if (ret != ACLNN_SUCCESS) {
		925	+ return ret;
		926	+ }
		927	+ // 校验 ori_topk_length 第三个维度
		928	+ ret = CheckTopkLengthSingleDimSmla(oriTopkLengthOptional, "ori_topk_length",
		929	+ SMLA_TOPK_LENGTH_THIRD_AXIS, numHeadsKv, "num_heads_kv",
		930	+ layoutQOptional);
		931	+ if (ret != ACLNN_SUCCESS) {
		932	+ return ret;
		933	+ }
		934	+ } else if (strcmp(layoutQOptional, "TND") == 0) {
		935	+ // 校验 ori_topk_length 第二个维度
		936	+ aclnnStatus ret = CheckTopkLengthSingleDimSmla(oriTopkLengthOptional, "ori_topk_length",
		937	+ SMLA_TOPK_LENGTH_SECOND_AXIS, numHeadsKv,
		938	+ "num_heads_kv", layoutQOptional);
		939	+ if (ret != ACLNN_SUCCESS) {
		940	+ return ret;
		941	+ }
		942	+ }
		943	+ }
824	}	944	}
825	if (hasCmpKv) {	945	if (hasCmpKv) {
826	std::string cmpKvSource;	946	std::string cmpKvSource;
	@@ -881,6 +1001,41 @@ aclnnStatus CheckConsistencySmla(const aclTensor *cuSeqlensQOptional, const aclT
881	return ACLNN_ERR_PARAM_INVALID;	1001	return ACLNN_ERR_PARAM_INVALID;
882	}	1002	}
883	}	1003	}
		1004	+ // 校验 cmp_topk_length 维度一致性
		1005	+ if (cmpTopk != 0 &&
		1006	+ cmpMaskMode == static_cast<int64_t>(SparseModeSmla::DEFAULT_MASK) &&
		1007	+ IsTensorExistSmla(cmpTopkLengthOptional)) {
		1008	+ if (strcmp(layoutQOptional, "BSND") == 0) {
		1009	+ // 校验 cmp_topk_length 第一个维度
		1010	+ aclnnStatus ret = CheckTopkLengthFirstDimSmla(cmpTopkLengthOptional, "cmp_topk_length",
		1011	+ queryBatchSize, querySource);
		1012	+ if (ret != ACLNN_SUCCESS) {
		1013	+ return ret;
		1014	+ }
		1015	+ // 校验 cmp_topk_length 第二个维度
		1016	+ ret = CheckTopkLengthSingleDimSmla(cmpTopkLengthOptional, "cmp_topk_length",
		1017	+ SMLA_TOPK_LENGTH_SECOND_AXIS, maxSeqlenQ, "max_seqlen_q",
		1018	+ layoutQOptional);
		1019	+ if (ret != ACLNN_SUCCESS) {
		1020	+ return ret;
		1021	+ }
		1022	+ // 校验 cmp_topk_length 第三个维度
		1023	+ ret = CheckTopkLengthSingleDimSmla(cmpTopkLengthOptional, "cmp_topk_length",
		1024	+ SMLA_TOPK_LENGTH_THIRD_AXIS, numHeadsKv, "num_heads_kv",
		1025	+ layoutQOptional);
		1026	+ if (ret != ACLNN_SUCCESS) {
		1027	+ return ret;
		1028	+ }
		1029	+ } else if (strcmp(layoutQOptional, "TND") == 0) {
		1030	+ // 校验 cmp_topk_length 第二个维度
		1031	+ aclnnStatus ret = CheckTopkLengthSingleDimSmla(cmpTopkLengthOptional, "cmp_topk_length",
		1032	+ SMLA_TOPK_LENGTH_SECOND_AXIS, numHeadsKv,
		1033	+ "num_heads_kv", layoutQOptional);
		1034	+ if (ret != ACLNN_SUCCESS) {
		1035	+ return ret;
		1036	+ }
		1037	+ }
		1038	+ }
884	}	1039	}
885	return ACLNN_SUCCESS;	1040	return ACLNN_SUCCESS;
886	}	1041	}
	@@ -908,7 +1063,8 @@ static aclnnStatus ParamsCheck(const aclTensor *cuSeqlensQOptional, const aclTen
908	CheckConsistencySmla(cuSeqlensQOptional, cuSeqlensOriKvOptional, cuSeqlensCmpKvOptional, sequsedQOptional,	1063	CheckConsistencySmla(cuSeqlensQOptional, cuSeqlensOriKvOptional, cuSeqlensCmpKvOptional, sequsedQOptional,
909	sequsedOriKvOptional, sequsedCmpKvOptional, cmpResidualKvOptional, oriTopkLengthOptional,	1064	sequsedOriKvOptional, sequsedCmpKvOptional, cmpResidualKvOptional, oriTopkLengthOptional,
910	cmpTopkLengthOptional, batchSize, layoutQOptional, layoutKvOptional, hasOriKv, hasCmpKv,	1065	cmpTopkLengthOptional, batchSize, layoutQOptional, layoutKvOptional, hasOriKv, hasCmpKv,
911	- socVersion, metaData) == ACLNN_SUCCESS) {	1066	+ oriTopk, cmpTopk, oriMaskMode, cmpMaskMode, maxSeqlenQ, numHeadsKv, socVersion,
		1067	+ metaData) == ACLNN_SUCCESS) {
912	return ACLNN_SUCCESS;	1068	return ACLNN_SUCCESS;
913	} else {	1069	} else {
914	return ACLNN_ERR_PARAM_INVALID;	1070	return ACLNN_ERR_PARAM_INVALID;

Mattention/sparse_flash_mla_metadata/op_kernel_aicpu/sparse_flash_mla_metadata_aicpu.cpp+88-31

	@@ -92,13 +92,12 @@ bool SparseFlashMlaMetadataCpuKernel::ParamsCheck()
92	if (layoutQ_ == "TND") {	92	if (layoutQ_ == "TND") {
93	if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) {	93	if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) {
94	const int32_t cuSeqlensQPtr = static_cast<const int32_t >(cuSeqlensQ_->GetData());	94	const int32_t cuSeqlensQPtr = static_cast<const int32_t >(cuSeqlensQ_->GetData());
		95	+ // 校验 cu_seqlens_q 首元素为 0
		96	+ if (cuSeqlensQPtr[0] != 0) {
		97	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_q should be 0, but got %d", cuSeqlensQPtr[0]);
		98	+ return false;
		99	+ }
95	for (int i = 0; i < batchSize + 1; i++) {	100	for (int i = 0; i < batchSize + 1; i++) {
96	- // 校验 cu_seqlens_q 元素非负
97	- if (cuSeqlensQPtr[i] < 0) {
98	- KERNEL_LOG_ERROR("The elements in cu_seqlens_q should be >= 0, but got cu_seqlens_q[%d] = %d", i,
99	- cuSeqlensQPtr[i]);
100	- return false;
101	- }
102	// 校验 cu_seqlens_q 元素递增	101	// 校验 cu_seqlens_q 元素递增
103	if (i > 0 && cuSeqlensQPtr[i - 1] > cuSeqlensQPtr[i]) {	102	if (i > 0 && cuSeqlensQPtr[i - 1] > cuSeqlensQPtr[i]) {
104	KERNEL_LOG_ERROR("The elements in cu_seqlens_q must be in ascending order, "	103	KERNEL_LOG_ERROR("The elements in cu_seqlens_q must be in ascending order, "
	@@ -112,6 +111,9 @@ bool SparseFlashMlaMetadataCpuKernel::ParamsCheck()
112	// 校验 seqused_q 元素	111	// 校验 seqused_q 元素
113	if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) {	112	if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) {
114	const int32_t sequsedQPtr = static_cast<const int32_t >(sequsedQ_->GetData());	113	const int32_t sequsedQPtr = static_cast<const int32_t >(sequsedQ_->GetData());
		114	+ const int32_t *cuSeqlensQPtr = (layoutQ_ == "TND" && cuSeqlensQ_ != nullptr &&
		115	+ cuSeqlensQ_->GetData() != nullptr) ?
		116	+ static_cast<const int32_t *>(cuSeqlensQ_->GetData()) : nullptr;
115	for (int i = 0; i < batchSize; i++) {	117	for (int i = 0; i < batchSize; i++) {
116	// 校验 seqused_q 元素非负	118	// 校验 seqused_q 元素非负
117	if (sequsedQPtr[i] < 0) {	119	if (sequsedQPtr[i] < 0) {
	@@ -119,6 +121,20 @@ bool SparseFlashMlaMetadataCpuKernel::ParamsCheck()
119	sequsedQPtr[i]);	121	sequsedQPtr[i]);
120	return false;	122	return false;
121	}	123	}
		124	+ // 校验 seqused_q 元素不大于 max_seqlen_q (BSND) 或 cu_seqlens_q 序列长度 (TND)
		125	+ if (layoutQ_ == "BSND" && sequsedQPtr[i] > maxSeqlenQ_) {
		126	+ KERNEL_LOG_ERROR("The elements in seqused_q should not be greater than max_seqlen_q %d, "
		127	+ "but got seqused_q[%d] = %d", maxSeqlenQ_, i, sequsedQPtr[i]);
		128	+ return false;
		129	+ }
		130	+ if (cuSeqlensQPtr != nullptr) {
		131	+ int32_t seqLen = cuSeqlensQPtr[i + 1] - cuSeqlensQPtr[i];
		132	+ if (sequsedQPtr[i] > seqLen) {
		133	+ KERNEL_LOG_ERROR("The elements in seqused_q should not be greater than the sequence length "
		134	+ "from cu_seqlens_q %d, but got seqused_q[%d] = %d", seqLen, i, sequsedQPtr[i]);
		135	+ return false;
		136	+ }
		137	+ }
122	}	138	}
123	}	139	}
124	if (hasOriKv_) {	140	if (hasOriKv_) {
	@@ -126,14 +142,13 @@ bool SparseFlashMlaMetadataCpuKernel::ParamsCheck()
126	if (layoutKv_ == "TND") {	142	if (layoutKv_ == "TND") {
127	if (cuSeqlensOriKv_ != nullptr && cuSeqlensOriKv_->GetData() != nullptr) {	143	if (cuSeqlensOriKv_ != nullptr && cuSeqlensOriKv_->GetData() != nullptr) {
128	const int32_t cuSeqlensOriKvPtr = static_cast<const int32_t >(cuSeqlensOriKv_->GetData());	144	const int32_t cuSeqlensOriKvPtr = static_cast<const int32_t >(cuSeqlensOriKv_->GetData());
		145	+ // 校验 cu_seqlens_ori_kv 首元素为 0
		146	+ if (cuSeqlensOriKvPtr[0] != 0) {
		147	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_ori_kv should be 0, but got %d",
		148	+ cuSeqlensOriKvPtr[0]);
		149	+ return false;
		150	+ }
129	for (int i = 0; i < batchSize + 1; i++) {	151	for (int i = 0; i < batchSize + 1; i++) {
130	- // 校验 cu_seqlens_ori_kv 元素非负
131	- if (cuSeqlensOriKvPtr[i] < 0) {
132	- KERNEL_LOG_ERROR("The elements in cu_seqlens_ori_kv should be >= 0, "
133	- "but got cu_seqlens_ori_kv[%d] = %d",
134	- i, cuSeqlensOriKvPtr[i]);
135	- return false;
136	- }
137	// 校验 cu_seqlens_ori_kv 元素递增	152	// 校验 cu_seqlens_ori_kv 元素递增
138	if (i > 0 && cuSeqlensOriKvPtr[i - 1] > cuSeqlensOriKvPtr[i]) {	153	if (i > 0 && cuSeqlensOriKvPtr[i - 1] > cuSeqlensOriKvPtr[i]) {
139	KERNEL_LOG_ERROR("The elements in cu_seqlens_ori_kv must be in ascending order, "	154	KERNEL_LOG_ERROR("The elements in cu_seqlens_ori_kv must be in ascending order, "
	@@ -147,6 +162,10 @@ bool SparseFlashMlaMetadataCpuKernel::ParamsCheck()
147	// 校验 seqused_ori_kv 元素	162	// 校验 seqused_ori_kv 元素
148	if (sequsedOriKv_ != nullptr && sequsedOriKv_->GetData() != nullptr) {	163	if (sequsedOriKv_ != nullptr && sequsedOriKv_->GetData() != nullptr) {
149	const int32_t sequsedOriKvPtr = static_cast<const int32_t >(sequsedOriKv_->GetData());	164	const int32_t sequsedOriKvPtr = static_cast<const int32_t >(sequsedOriKv_->GetData());
		165	+ const int32_t *cuSeqlensOriKvPtr = (layoutKv_ == "TND" && cuSeqlensOriKv_ != nullptr &&
		166	+ cuSeqlensOriKv_->GetData() != nullptr) ?
		167	+ static_cast<const int32_t *>(cuSeqlensOriKv_->GetData()) :
		168	+ nullptr;
150	for (int i = 0; i < batchSize; i++) {	169	for (int i = 0; i < batchSize; i++) {
151	// 校验 seqused_ori_kv 元素非负	170	// 校验 seqused_ori_kv 元素非负
152	if (sequsedOriKvPtr[i] < 0) {	171	if (sequsedOriKvPtr[i] < 0) {
	@@ -154,10 +173,27 @@ bool SparseFlashMlaMetadataCpuKernel::ParamsCheck()
154	i, sequsedOriKvPtr[i]);	173	i, sequsedOriKvPtr[i]);
155	return false;	174	return false;
156	}	175	}
		176	+ // 校验 seqused_ori_kv 元素不大于 max_seqlen_ori_kv (BSND) 或 cu_seqlens_ori_kv 序列长度 (TND)
		177	+ if (layoutKv_ == "BSND" && sequsedOriKvPtr[i] > maxSeqlenOriKv_) {
		178	+ KERNEL_LOG_ERROR("The elements in seqused_ori_kv should not be greater than "
		179	+ "max_seqlen_ori_kv %d, but got seqused_ori_kv[%d] = %d",
		180	+ maxSeqlenOriKv_, i, sequsedOriKvPtr[i]);
		181	+ return false;
		182	+ }
		183	+ if (cuSeqlensOriKvPtr != nullptr) {
		184	+ int32_t seqLen = cuSeqlensOriKvPtr[i + 1] - cuSeqlensOriKvPtr[i];
		185	+ if (sequsedOriKvPtr[i] > seqLen) {
		186	+ KERNEL_LOG_ERROR("The elements in seqused_ori_kv should not be greater than the sequence "
		187	+ "length from cu_seqlens_ori_kv %d, but got seqused_ori_kv[%d] = %d",
		188	+ seqLen, i, sequsedOriKvPtr[i]);
		189	+ return false;
		190	+ }
		191	+ }
157	}	192	}
158	}	193	}
159	// 校验 ori_topk_length 元素	194	// 校验 ori_topk_length 元素
160	- if (oriTopkLength_ != nullptr && oriTopkLength_->GetData() != nullptr) {	195	+ if (oriTopK_ != 0 && oriMaskMode_ == static_cast<int32_t>(SparseMode::DEFAULT_MASK) &&
		196	+ oriTopkLength_ != nullptr && oriTopkLength_->GetData() != nullptr) {
161	// 校验 ori_topk_length 元素数量	197	// 校验 ori_topk_length 元素数量
162	int32_t sumOfQuerySeq = GetSumOfQuerySeq();	198	int32_t sumOfQuerySeq = GetSumOfQuerySeq();
163	const int32_t oriTopkLengthPtr = static_cast<const int32_t >(oriTopkLength_->GetData());	199	const int32_t oriTopkLengthPtr = static_cast<const int32_t >(oriTopkLength_->GetData());
	@@ -187,14 +223,13 @@ bool SparseFlashMlaMetadataCpuKernel::ParamsCheck()
187	// 校验 cu_seqlens_cmp_kv 元素	223	// 校验 cu_seqlens_cmp_kv 元素
188	if (cuSeqlensCmpKv_ != nullptr && cuSeqlensCmpKv_->GetData() != nullptr) {	224	if (cuSeqlensCmpKv_ != nullptr && cuSeqlensCmpKv_->GetData() != nullptr) {
189	const int32_t cuSeqlensCmpKvPtr = static_cast<const int32_t >(cuSeqlensCmpKv_->GetData());	225	const int32_t cuSeqlensCmpKvPtr = static_cast<const int32_t >(cuSeqlensCmpKv_->GetData());
		226	+ // 校验 cu_seqlens_cmp_kv 首元素为 0
		227	+ if (cuSeqlensCmpKvPtr[0] != 0) {
		228	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_cmp_kv should be 0, but got %d",
		229	+ cuSeqlensCmpKvPtr[0]);
		230	+ return false;
		231	+ }
190	for (int i = 0; i < batchSize + 1; i++) {	232	for (int i = 0; i < batchSize + 1; i++) {
191	- // 校验 cu_seqlens_cmp_kv 元素非负
192	- if (cuSeqlensCmpKvPtr[i] < 0) {
193	- KERNEL_LOG_ERROR("The elements in cu_seqlens_cmp_kv should be >= 0, "
194	- "but got cu_seqlens_cmp_kv[%d] = %d",
195	- i, cuSeqlensCmpKvPtr[i]);
196	- return false;
197	- }
198	// 校验 cu_seqlens_cmp_kv 元素递增	233	// 校验 cu_seqlens_cmp_kv 元素递增
199	if (i > 0 && cuSeqlensCmpKvPtr[i - 1] > cuSeqlensCmpKvPtr[i]) {	234	if (i > 0 && cuSeqlensCmpKvPtr[i - 1] > cuSeqlensCmpKvPtr[i]) {
200	KERNEL_LOG_ERROR("The elements in cu_seqlens_cmp_kv must be in ascending order, "	235	KERNEL_LOG_ERROR("The elements in cu_seqlens_cmp_kv must be in ascending order, "
	@@ -208,6 +243,10 @@ bool SparseFlashMlaMetadataCpuKernel::ParamsCheck()
208	// 校验 seqused_cmp_kv 元素	243	// 校验 seqused_cmp_kv 元素
209	if (sequsedCmpKv_ != nullptr && sequsedCmpKv_->GetData() != nullptr) {	244	if (sequsedCmpKv_ != nullptr && sequsedCmpKv_->GetData() != nullptr) {
210	const int32_t sequsedCmpKvPtr = static_cast<const int32_t >(sequsedCmpKv_->GetData());	245	const int32_t sequsedCmpKvPtr = static_cast<const int32_t >(sequsedCmpKv_->GetData());
		246	+ const int32_t *cuSeqlensCmpKvPtr = (layoutKv_ == "TND" && cuSeqlensCmpKv_ != nullptr &&
		247	+ cuSeqlensCmpKv_->GetData() != nullptr) ?
		248	+ static_cast<const int32_t *>(cuSeqlensCmpKv_->GetData()) :
		249	+ nullptr;
211	for (int i = 0; i < batchSize; i++) {	250	for (int i = 0; i < batchSize; i++) {
212	// 校验 seqused_cmp_kv 元素非负	251	// 校验 seqused_cmp_kv 元素非负
213	if (sequsedCmpKvPtr[i] < 0) {	252	if (sequsedCmpKvPtr[i] < 0) {
	@@ -215,23 +254,39 @@ bool SparseFlashMlaMetadataCpuKernel::ParamsCheck()
215	i, sequsedCmpKvPtr[i]);	254	i, sequsedCmpKvPtr[i]);
216	return false;	255	return false;
217	}	256	}
		257	+ // 校验 seqused_cmp_kv 元素不大于 max_seqlen_cmp_kv (BSND) 或 cu_seqlens_cmp_kv 序列长度 (TND)
		258	+ if (layoutKv_ == "BSND" && sequsedCmpKvPtr[i] > maxSeqlenCmpKv_) {
		259	+ KERNEL_LOG_ERROR("The elements in seqused_cmp_kv should not be greater than "
		260	+ "max_seqlen_cmp_kv %d, but got seqused_cmp_kv[%d] = %d",
		261	+ maxSeqlenCmpKv_, i, sequsedCmpKvPtr[i]);
		262	+ return false;
		263	+ }
		264	+ if (cuSeqlensCmpKvPtr != nullptr) {
		265	+ int32_t seqLen = cuSeqlensCmpKvPtr[i + 1] - cuSeqlensCmpKvPtr[i];
		266	+ if (sequsedCmpKvPtr[i] > seqLen) {
		267	+ KERNEL_LOG_ERROR("The elements in seqused_cmp_kv should not be greater than the sequence "
		268	+ "length from cu_seqlens_cmp_kv %d, but got seqused_cmp_kv[%d] = %d",
		269	+ seqLen, i, sequsedCmpKvPtr[i]);
		270	+ return false;
		271	+ }
		272	+ }
218	}	273	}
219	}	274	}
220	// 校验 cmp_residual_kv 元素	275	// 校验 cmp_residual_kv 元素
221	if (cmpResidualKv_ != nullptr && cmpResidualKv_->GetData() != nullptr) {	276	if (cmpResidualKv_ != nullptr && cmpResidualKv_->GetData() != nullptr) {
222	const int32_t cmpResidualKvPtr = static_cast<const int32_t >(cmpResidualKv_->GetData());	277	const int32_t cmpResidualKvPtr = static_cast<const int32_t >(cmpResidualKv_->GetData());
223	for (int i = 0; i < batchSize; i++) {	278	for (int i = 0; i < batchSize; i++) {
224	- // 校验 cmp_residual_kv 元素非负
225	if (cmpResidualKvPtr[i] < 0 \|\| cmpResidualKvPtr[i] >= cmpRatio_) {	279	if (cmpResidualKvPtr[i] < 0 \|\| cmpResidualKvPtr[i] >= cmpRatio_) {
226	- KERNEL_LOG_ERROR("The elements in cmp_residual_kv should be in [0, cmpRatio_), but got "	280	+ KERNEL_LOG_ERROR("The elements in cmp_residual_kv should be in [0, cmpRatio_(%d)), but got "
227	- "cmp_residual_kv[%d] = %d",	281	+ "cmp_residual_kv[%d] = %d", cmpRatio_,
228	i, cmpResidualKvPtr[i]);	282	i, cmpResidualKvPtr[i]);
229	return false;	283	return false;
230	}	284	}
231	}	285	}
232	}	286	}
233	// 校验 cmp_topk_length 元素	287	// 校验 cmp_topk_length 元素
234	- if (cmpTopkLength_ != nullptr && cmpTopkLength_->GetData() != nullptr) {	288	+ if (cmpTopK_ != 0 && cmpMaskMode_ == static_cast<int32_t>(SparseMode::DEFAULT_MASK) &&
		289	+ cmpTopkLength_ != nullptr && cmpTopkLength_->GetData() != nullptr) {
235	// 校验 cmp_topk_length 元素数量	290	// 校验 cmp_topk_length 元素数量
236	int32_t sumOfQuerySeq = GetSumOfQuerySeq();	291	int32_t sumOfQuerySeq = GetSumOfQuerySeq();
237	const int32_t cmpTopkLengthPtr = static_cast<const int32_t >(cmpTopkLength_->GetData());	292	const int32_t cmpTopkLengthPtr = static_cast<const int32_t >(cmpTopkLength_->GetData());
	@@ -419,7 +474,8 @@ uint32_t SparseFlashMlaMetadataCpuKernel::GetBsStride(uint32_t bIdx, uint32_t s1
419	uint32_t SparseFlashMlaMetadataCpuKernel::GetOriTopkLength(uint32_t bsStride)	474	uint32_t SparseFlashMlaMetadataCpuKernel::GetOriTopkLength(uint32_t bsStride)
420	{	475	{
421	// 尝试使用 oriTopkLength_	476	// 尝试使用 oriTopkLength_
422	- if (oriTopkLength_ != nullptr && oriTopkLength_->GetData() != nullptr) {	477	+ if (oriTopK_ != 0 && oriMaskMode_ == static_cast<int32_t>(SparseMode::DEFAULT_MASK) &&
		478	+ oriTopkLength_ != nullptr && oriTopkLength_->GetData() != nullptr) {
423	const int32_t oriTopkPtr = static_cast<const int32_t >(oriTopkLength_->GetData());	479	const int32_t oriTopkPtr = static_cast<const int32_t >(oriTopkLength_->GetData());
424	return static_cast<uint32_t>(oriTopkPtr[bsStride]);	480	return static_cast<uint32_t>(oriTopkPtr[bsStride]);
425	}	481	}
	@@ -430,7 +486,8 @@ uint32_t SparseFlashMlaMetadataCpuKernel::GetOriTopkLength(uint32_t bsStride)
430	uint32_t SparseFlashMlaMetadataCpuKernel::GetCmpTopkLength(uint32_t bsStride)	486	uint32_t SparseFlashMlaMetadataCpuKernel::GetCmpTopkLength(uint32_t bsStride)
431	{	487	{
432	// 尝试使用 cmpTopkLength_	488	// 尝试使用 cmpTopkLength_
433	- if (cmpTopkLength_ != nullptr && cmpTopkLength_->GetData() != nullptr) {	489	+ if (cmpTopK_ != 0 && cmpMaskMode_ == static_cast<int32_t>(SparseMode::DEFAULT_MASK) &&
		490	+ cmpTopkLength_ != nullptr && cmpTopkLength_->GetData() != nullptr) {
434	const int32_t cmpTopkPtr = static_cast<const int32_t >(cmpTopkLength_->GetData());	491	const int32_t cmpTopkPtr = static_cast<const int32_t >(cmpTopkLength_->GetData());
435	return static_cast<uint32_t>(cmpTopkPtr[bsStride]);	492	return static_cast<uint32_t>(cmpTopkPtr[bsStride]);
436	}	493	}
	@@ -472,8 +529,8 @@ uint32_t SparseFlashMlaMetadataCpuKernel::GetOriS2SeqSize(uint32_t bIdx)
472	return static_cast<uint32_t>(s2Ptr[bIdx + 1U] - s2Ptr[bIdx]);	529	return static_cast<uint32_t>(s2Ptr[bIdx + 1U] - s2Ptr[bIdx]);
473	}	530	}
474	}	531	}
475	- // 如果 max_seqlen_ori_kv 没传入，且 ori_kv 为稀疏的，则尝试从 topk 中获取	532	+ // 如果是PA场景，或 max_seqlen_ori_kv 没传入，且 ori_kv 为稀疏的，则尝试从 topk 中获取
476	- if (maxSeqlenOriKv_ == 0 && isSparseOriKv_) {	533	+ if ((layoutKv_ == "PA_BBND" \|\| maxSeqlenOriKv_ == 0) && isSparseOriKv_) {
477	return UINT32_MAX;	534	return UINT32_MAX;
478	}	535	}
479	// 使用 max_seqlen_ori_kv	536	// 使用 max_seqlen_ori_kv
	@@ -495,8 +552,8 @@ uint32_t SparseFlashMlaMetadataCpuKernel::GetCmpS2SeqSize(uint32_t bIdx)
495	return static_cast<uint32_t>(s2Ptr[bIdx + 1U] - s2Ptr[bIdx]);	552	return static_cast<uint32_t>(s2Ptr[bIdx + 1U] - s2Ptr[bIdx]);
496	}	553	}
497	}	554	}
498	- // 如果 max_seqlen_cmp_kv 没传入，且 cmp_kv 为稀疏的，则尝试从topk中获取	555	+ // 如果是PA场景，或 max_seqlen_cmp_kv 没传入，且 cmp_kv 为稀疏的，则尝试从topk中获取
499	- if (maxSeqlenCmpKv_ == 0 && isSparseCmpKv_) {	556	+ if ((layoutKv_ == "PA_BBND" \|\| maxSeqlenCmpKv_ == 0) && isSparseCmpKv_) {
500	return UINT32_MAX;	557	return UINT32_MAX;
501	}	558	}
502	// 使用 max_seqlen_cmp_kv	559	// 使用 max_seqlen_cmp_kv

Mattention/sparse_lightning_indexer_kl_loss_grad_metadata/op_host/sparse_lightning_indexer_kl_loss_grad_metadata_check.h+7-0

	@@ -204,6 +204,13 @@ aclnnStatus CheckSingleParamSli(int64_t batchSize, int64_t maxSeqlenQ, int64_t m
204	"The value of layout_q must be equal to that of layout_k");	204	"The value of layout_q must be equal to that of layout_k");
205	return ACLNN_ERR_PARAM_INVALID;	205	return ACLNN_ERR_PARAM_INVALID;
206	}	206	}
		207	+ // 校验 layout_q 为 BSND 时，max_seqlen_q 必须大于 0
		208	+ if (strcmp(layoutQOptional, "BSND") == 0 && maxSeqlenQ <= 0) {
		209	+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(SLI_ACLNN_OP_NAME, "max_seqlen_q", std::to_string(maxSeqlenQ),
		210	+ "When layout_q is BSND, the value of max_seqlen_q "
		211	+ "must be equal to the size of the second axis of q");
		212	+ return ACLNN_ERR_PARAM_INVALID;
		213	+ }
207	// 核心数校验	214	// 核心数校验
208	if (aicCoreNum == 0) {	215	if (aicCoreNum == 0) {
209	OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(SLI_ACLNN_OP_NAME, "aic_core_num", std::to_string(aicCoreNum),	216	OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(SLI_ACLNN_OP_NAME, "aic_core_num", std::to_string(aicCoreNum),

Mattention/sparse_lightning_indexer_kl_loss_grad_metadata/op_kernel_aicpu/arch35/sparse_lightning_indexer_kl_loss_grad_metadata_aicpu_arch35.h+27-12

	@@ -167,13 +167,12 @@ inline bool SparseLightningIndexerKLLossGradMetadataCpuKernelArch35::ParamsCheck
167	if (layoutQ_ == "TND") {	167	if (layoutQ_ == "TND") {
168	if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) {	168	if (cuSeqlensQ_ != nullptr && cuSeqlensQ_->GetData() != nullptr) {
169	const int32_t cuSeqlensQPtr = static_cast<const int32_t>(cuSeqlensQ_->GetData());	169	const int32_t cuSeqlensQPtr = static_cast<const int32_t>(cuSeqlensQ_->GetData());
		170	+ // 校验 cu_seqlens_q 首元素为 0
		171	+ if (cuSeqlensQPtr[0] != 0) {
		172	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_q should be 0, but got %d", cuSeqlensQPtr[0]);
		173	+ return false;
		174	+ }
170	for (int i = 0; i < batchSize + 1; i++) {	175	for (int i = 0; i < batchSize + 1; i++) {
171	- // 校验 cu_seqlens_q 元素非负
172	- if (cuSeqlensQPtr[i] < 0) {
173	- KERNEL_LOG_ERROR("The elements in cu_seqlens_q should be >= 0, but got cu_seqlens_q[%d] = %d",
174	- i, cuSeqlensQPtr[i]);
175	- return false;
176	- }
177	// 校验 cu_seqlens_q 元素递增	176	// 校验 cu_seqlens_q 元素递增
178	if (i > 0 && cuSeqlensQPtr[i - 1] > cuSeqlensQPtr[i]) {	177	if (i > 0 && cuSeqlensQPtr[i - 1] > cuSeqlensQPtr[i]) {
179	KERNEL_LOG_ERROR("The elements in cu_seqlens_q must be in ascending order, "	178	KERNEL_LOG_ERROR("The elements in cu_seqlens_q must be in ascending order, "
	@@ -188,13 +187,12 @@ inline bool SparseLightningIndexerKLLossGradMetadataCpuKernelArch35::ParamsCheck
188	if (layoutK_ == "TND") {	187	if (layoutK_ == "TND") {
189	if (cuSeqlensK_ != nullptr && cuSeqlensK_->GetData() != nullptr) {	188	if (cuSeqlensK_ != nullptr && cuSeqlensK_->GetData() != nullptr) {
190	const int32_t cuSeqlensKPtr = static_cast<const int32_t>(cuSeqlensK_->GetData());	189	const int32_t cuSeqlensKPtr = static_cast<const int32_t>(cuSeqlensK_->GetData());
		190	+ // 校验 cu_seqlens_k 首元素为 0
		191	+ if (cuSeqlensKPtr[0] != 0) {
		192	+ KERNEL_LOG_ERROR("The first element of cu_seqlens_k should be 0, but got %d", cuSeqlensKPtr[0]);
		193	+ return false;
		194	+ }
191	for (int i = 0; i < batchSize + 1; i++) {	195	for (int i = 0; i < batchSize + 1; i++) {
192	- // 校验 cu_seqlens_k 元素非负
193	- if (cuSeqlensKPtr[i] < 0) {
194	- KERNEL_LOG_ERROR("The elements in cu_seqlens_k should be >= 0, but got cu_seqlens_k[%d] = %d",
195	- i, cuSeqlensKPtr[i]);
196	- return false;
197	- }
198	// 校验 cu_seqlens_k 元素递增	196	// 校验 cu_seqlens_k 元素递增
199	if (i > 0 && cuSeqlensKPtr[i - 1] > cuSeqlensKPtr[i]) {	197	if (i > 0 && cuSeqlensKPtr[i - 1] > cuSeqlensKPtr[i]) {
200	KERNEL_LOG_ERROR("The elements in cu_seqlens_k must be in ascending order, "	198	KERNEL_LOG_ERROR("The elements in cu_seqlens_k must be in ascending order, "
	@@ -208,12 +206,29 @@ inline bool SparseLightningIndexerKLLossGradMetadataCpuKernelArch35::ParamsCheck
208	// 校验 seqused_q 元素非负	206	// 校验 seqused_q 元素非负
209	if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) {	207	if (sequsedQ_ != nullptr && sequsedQ_->GetData() != nullptr) {
210	const int32_t sequsedQPtr = static_cast<const int32_t>(sequsedQ_->GetData());	208	const int32_t sequsedQPtr = static_cast<const int32_t>(sequsedQ_->GetData());
		209	+ const int32_t *cuSeqlensQPtr = (layoutQ_ == "TND" && cuSeqlensQ_ != nullptr &&
		210	+ cuSeqlensQ_->GetData() != nullptr) ?
		211	+ static_cast<const int32_t*>(cuSeqlensQ_->GetData()) : nullptr;
211	for (int i = 0; i < batchSize; i++) {	212	for (int i = 0; i < batchSize; i++) {
212	if (sequsedQPtr[i] < 0) {	213	if (sequsedQPtr[i] < 0) {
213	KERNEL_LOG_ERROR("The elements in seqused_q should be >= 0, but got seqused_q[%d] = %d",	214	KERNEL_LOG_ERROR("The elements in seqused_q should be >= 0, but got seqused_q[%d] = %d",
214	i, sequsedQPtr[i]);	215	i, sequsedQPtr[i]);
215	return false;	216	return false;
216	}	217	}
		218	+ // 校验 seqused_q 元素不大于 max_seqlen_q (BSND) 或 cu_seqlens_q 序列长度 (TND)
		219	+ if (layoutQ_ == "BSND" && sequsedQPtr[i] > maxSeqlenQ_) {
		220	+ KERNEL_LOG_ERROR("The elements in seqused_q should not be greater than max_seqlen_q %d, "
		221	+ "but got seqused_q[%d] = %d", maxSeqlenQ_, i, sequsedQPtr[i]);
		222	+ return false;
		223	+ }
		224	+ if (cuSeqlensQPtr != nullptr) {
		225	+ int32_t seqLen = cuSeqlensQPtr[i + 1] - cuSeqlensQPtr[i];
		226	+ if (sequsedQPtr[i] > seqLen) {
		227	+ KERNEL_LOG_ERROR("The elements in seqused_q should not be greater than the sequence length "
		228	+ "from cu_seqlens_q %d, but got seqused_q[%d] = %d", seqLen, i, sequsedQPtr[i]);
		229	+ return false;
		230	+ }
		231	+ }
217	}	232	}
218	}	233	}
219	// 校验 seqused_k 元素非负	234	// 校验 seqused_k 元素非负