支持 SparseFlashMla 与 SparseFlashMlaMetadata 算子功能增强-ops-transformer-AtomGit

Mattention/sparse_flash_mla/examples/test_aclnn_sparse_flash_mla.cpp+4-3

	@@ -293,7 +293,7 @@ int main()
293	CHECK_RET(ret == ACL_SUCCESS, return ret);	293	CHECK_RET(ret == ACL_SUCCESS, return ret);
294		294
295	char layoutQ[] = "TND";	295	char layoutQ[] = "TND";
296	- char layoutKv[] = "PA_BNBD";	296	+ char layoutKv[] = "PA_BBND";
297		297
298	uint64_t metadataWorkspaceSize = 0;	298	uint64_t metadataWorkspaceSize = 0;
299	aclOpExecutor* metadataExecutor = nullptr;	299	aclOpExecutor* metadataExecutor = nullptr;
	@@ -302,7 +302,8 @@ int main()
302	ret = aclnnSparseFlashMlaMetadataGetWorkspaceSize(	302	ret = aclnnSparseFlashMlaMetadataGetWorkspaceSize(
303	cuSeqLensQ, cuSeqLensOriKv, cuSeqLensCmpKv,	303	cuSeqLensQ, cuSeqLensOriKv, cuSeqLensCmpKv,
304	seqUsedQ, seqUsedOriKv,	304	seqUsedQ, seqUsedOriKv,
305	- N1, N2, D, B, S1, S2,	305	+ nullptr, nullptr, nullptr, nullptr,
		306	+ N1, N2, D, B, S1, S2, cmpKvLen,
306	0, K, cmpRatio,	307	0, K, cmpRatio,
307	oriMaskMode, cmpMaskMode,	308	oriMaskMode, cmpMaskMode,
308	oriWinLeft, oriWinRight,	309	oriWinLeft, oriWinRight,
	@@ -341,7 +342,7 @@ int main()
341	oriMaskMode, cmpMaskMode,	342	oriMaskMode, cmpMaskMode,
342	oriWinLeft, oriWinRight,	343	oriWinLeft, oriWinRight,
343	layoutQ, layoutKv,	344	layoutQ, layoutKv,
344	- 1, 0, 0,	345	+ 1,
345	false,	346	false,
346	attnOut, softmaxLse,	347	attnOut, softmaxLse,
347	&workspaceSize, &executor);	348	&workspaceSize, &executor);

Mattention/sparse_flash_mla/op_host/sparse_flash_mla_def.cpp+4-6

	@@ -118,17 +118,15 @@ public:
118	.ParamType(OPTIONAL)	118	.ParamType(OPTIONAL)
119	.DataTypeList({ge::DT_FLOAT})	119	.DataTypeList({ge::DT_FLOAT})
120	.FormatList({ge::FORMAT_ND});	120	.FormatList({ge::FORMAT_ND});
121	- this->Attr("softmax_scale").AttrType(REQUIRED).Float(1.0);	121	+ this->Attr("softmax_scale").AttrType(OPTIONAL).Float(1.0);
122	- this->Attr("cmp_ratio").AttrType(REQUIRED).Int(1);	122	+ this->Attr("cmp_ratio").AttrType(OPTIONAL).Int(1);
123	- this->Attr("ori_mask_mode").AttrType(REQUIRED).Int(0); // ori_mask_mode默认值0	123	+ this->Attr("ori_mask_mode").AttrType(OPTIONAL).Int(0); // ori_mask_mode默认值0
124	- this->Attr("cmp_mask_mode").AttrType(REQUIRED).Int(0); // cmp_mask_mode默认值0	124	+ this->Attr("cmp_mask_mode").AttrType(OPTIONAL).Int(0); // cmp_mask_mode默认值0
125	this->Attr("ori_win_left").AttrType(OPTIONAL).Int(-1); // ori_win_left默认值-1	125	this->Attr("ori_win_left").AttrType(OPTIONAL).Int(-1); // ori_win_left默认值-1
126	this->Attr("ori_win_right").AttrType(OPTIONAL).Int(-1); // ori_win_right默认值-1	126	this->Attr("ori_win_right").AttrType(OPTIONAL).Int(-1); // ori_win_right默认值-1
127	this->Attr("layout_q").AttrType(OPTIONAL).String("BSND");	127	this->Attr("layout_q").AttrType(OPTIONAL).String("BSND");
128	this->Attr("layout_kv").AttrType(OPTIONAL).String("BSND");	128	this->Attr("layout_kv").AttrType(OPTIONAL).String("BSND");
129	this->Attr("topk_value_mode").AttrType(OPTIONAL).Int(1);	129	this->Attr("topk_value_mode").AttrType(OPTIONAL).Int(1);
130	- this->Attr("ori_kv_stride").AttrType(OPTIONAL).Int(0);
131	- this->Attr("cmp_kv_stride").AttrType(OPTIONAL).Int(0);
132	this->Attr("return_softmax_lse").AttrType(OPTIONAL).Bool(false);	130	this->Attr("return_softmax_lse").AttrType(OPTIONAL).Bool(false);
133		131
134	OpAICoreConfig aicore_config;	132	OpAICoreConfig aicore_config;

Mattention/sparse_flash_mla/op_host/sparse_flash_mla_tiling.cpp+181-133

	@@ -36,7 +36,17 @@ static const std::string CMP_BLOCK_TABLE_NAME = "cmp_block_table";
36	static const std::string SINKS_NAME = "sinks";	36	static const std::string SINKS_NAME = "sinks";
37	static const std::string METADATA_NAME = "metadata";	37	static const std::string METADATA_NAME = "metadata";
38	static const std::string ATTEN_OUT_NAME = "attn_out";	38	static const std::string ATTEN_OUT_NAME = "attn_out";
39	-static const std::map<std::string, std::vector<ge::DataType>> DTYPE_SUPPORT_MAP = {	39	+static bool IsNonEmptyOptionalTensor(const gert::Tensor *tensor)
		40	+{
		41	+ return tensor != nullptr && tensor->GetShapeSize() > 0;
		42	+}
		43	+
		44	+static bool IsPowerOfTwoInRange(uint32_t value, uint32_t minValue, uint32_t maxValue)
		45	+{
		46	+ return value >= minValue && value <= maxValue && (value & (value - 1U)) == 0U;
		47	+}
		48	+
		49	+static const std::map<std::string, std::vector<ge::DataType>> DTYPE_SUPPORT_MAP = {
40	{QUERY_NAME, {ge::DT_FLOAT16, ge::DT_BF16}},	50	{QUERY_NAME, {ge::DT_FLOAT16, ge::DT_BF16}},
41	{ORI_KV_NAME, {ge::DT_FLOAT16, ge::DT_BF16}},	51	{ORI_KV_NAME, {ge::DT_FLOAT16, ge::DT_BF16}},
42	{CMP_KV_NAME, {ge::DT_FLOAT16, ge::DT_BF16}},	52	{CMP_KV_NAME, {ge::DT_FLOAT16, ge::DT_BF16}},
	@@ -53,8 +63,8 @@ static const std::map<std::string, std::vector<ge::DataType>> DTYPE_SUPPORT_MAP
53		63
54	static const std::map<std::string, std::vector<SMLALayout>> LAYOUT_SUPPORT_MAP = {	64	static const std::map<std::string, std::vector<SMLALayout>> LAYOUT_SUPPORT_MAP = {
55	{QUERY_NAME, {SMLALayout::BSND, SMLALayout::TND}},	65	{QUERY_NAME, {SMLALayout::BSND, SMLALayout::TND}},
56	- {ORI_KV_NAME, {SMLALayout::PA_BBND, SMLALayout::TND, SMLALayout::BSND}},	66	+ {ORI_KV_NAME, {SMLALayout::PA_BBND, SMLALayout::TND, SMLALayout::BSND}},
57	- {CMP_KV_NAME, {SMLALayout::PA_BBND, SMLALayout::TND, SMLALayout::BSND}},	67	+ {CMP_KV_NAME, {SMLALayout::PA_BBND, SMLALayout::TND, SMLALayout::BSND}},
58	{ATTEN_OUT_NAME, {SMLALayout::BSND, SMLALayout::TND}},	68	{ATTEN_OUT_NAME, {SMLALayout::BSND, SMLALayout::TND}},
59	{ORI_SPARSE_INDICES, {SMLALayout::BSND, SMLALayout::TND}},	69	{ORI_SPARSE_INDICES, {SMLALayout::BSND, SMLALayout::TND}},
60	{CMP_SPARSE_INDICES, {SMLALayout::BSND, SMLALayout::TND}},	70	{CMP_SPARSE_INDICES, {SMLALayout::BSND, SMLALayout::TND}},
	@@ -101,7 +111,7 @@ std::string SMLALayoutToSerialString(SMLALayout layout)
101	switch (layout) {	111	switch (layout) {
102	case SMLALayout::BSND: return "BSND";	112	case SMLALayout::BSND: return "BSND";
103	case SMLALayout::TND: return "TND";	113	case SMLALayout::TND: return "TND";
104	- case SMLALayout::PA_BBND: return "PA_BBND";	114	+ case SMLALayout::PA_BBND: return "PA_BBND";
105	default: return "UNKNOWN";	115	default: return "UNKNOWN";
106	}	116	}
107	}	117	}
	@@ -113,13 +123,13 @@ struct SMLACompileInfo {
113	static const std::map<SMLALayout, std::vector<SMLAAxis>> SMLA_LAYOUT_AXIS_MAP = {	123	static const std::map<SMLALayout, std::vector<SMLAAxis>> SMLA_LAYOUT_AXIS_MAP = {
114	{SMLALayout::BSND, {SMLAAxis::B, SMLAAxis::S, SMLAAxis::N, SMLAAxis::D}},	124	{SMLALayout::BSND, {SMLAAxis::B, SMLAAxis::S, SMLAAxis::N, SMLAAxis::D}},
115	{SMLALayout::TND, {SMLAAxis::T, SMLAAxis::N, SMLAAxis::D}},	125	{SMLALayout::TND, {SMLAAxis::T, SMLAAxis::N, SMLAAxis::D}},
116	- {SMLALayout::PA_BBND, {SMLAAxis::Bn, SMLAAxis::Bs, SMLAAxis::N, SMLAAxis::D}},	126	+ {SMLALayout::PA_BBND, {SMLAAxis::Bn, SMLAAxis::Bs, SMLAAxis::N, SMLAAxis::D}},
117	};	127	};
118		128
119	static const std::map<SMLALayout, size_t> SMLA_LAYOUT_DIM_MAP = {	129	static const std::map<SMLALayout, size_t> SMLA_LAYOUT_DIM_MAP = {
120	{SMLALayout::BSND, DIM_NUM_FOUR},	130	{SMLALayout::BSND, DIM_NUM_FOUR},
121	{SMLALayout::TND, DIM_NUM_THREE},	131	{SMLALayout::TND, DIM_NUM_THREE},
122	- {SMLALayout::PA_BBND, DIM_NUM_FOUR},	132	+ {SMLALayout::PA_BBND, DIM_NUM_FOUR},
123	};	133	};
124		134
125	static std::string SMLADataTypeToSerialString(ge::DataType type)	135	static std::string SMLADataTypeToSerialString(ge::DataType type)
	@@ -142,9 +152,9 @@ ge::graphStatus SMLAInfoParser::CheckRequiredInOutExistence() const
142	return ge::GRAPH_FAILED);	152	return ge::GRAPH_FAILED);
143	OP_CHECK_IF(opParamInfo_.oriKv.tensor == nullptr, OP_LOGE(opName_, "tensor of oriKv is nullptr"),	153	OP_CHECK_IF(opParamInfo_.oriKv.tensor == nullptr, OP_LOGE(opName_, "tensor of oriKv is nullptr"),
144	return ge::GRAPH_FAILED);	154	return ge::GRAPH_FAILED);
145	- if (opParamInfo_.layoutKv == "PA_BBND") {	155	+ if (std::string(opParamInfo_.layoutKv) == "PA_BBND") {
146	OP_CHECK_IF(opParamInfo_.oriBlockTable.tensor == nullptr,	156	OP_CHECK_IF(opParamInfo_.oriBlockTable.tensor == nullptr,
147	- OP_LOGE(opName_, "tensor of oriBlockTable is nullptr when layoutKv is PA_BBND"),	157	+ OP_LOGE(opName_, "tensor of oriBlockTable is nullptr when layoutKv is PA_BBND"),
148	return ge::GRAPH_FAILED);	158	return ge::GRAPH_FAILED);
149	}	159	}
150	if (perfMode_ == SMLATemplateMode::CFA_TEMPLATE_MODE) {	160	if (perfMode_ == SMLATemplateMode::CFA_TEMPLATE_MODE) {
	@@ -270,8 +280,6 @@ ge::graphStatus SMLAInfoParser::GetAttrParaInfo()
270	opParamInfo_.cmpRatio = attrs->GetAttrPointer<uint32_t>(ATTR_CMP_RATIO_INDEX);	280	opParamInfo_.cmpRatio = attrs->GetAttrPointer<uint32_t>(ATTR_CMP_RATIO_INDEX);
271	opParamInfo_.oriMaskMode = attrs->GetAttrPointer<uint32_t>(ATTR_ORI_MASK_MODE_INDEX);	281	opParamInfo_.oriMaskMode = attrs->GetAttrPointer<uint32_t>(ATTR_ORI_MASK_MODE_INDEX);
272	opParamInfo_.cmpMaskMode = attrs->GetAttrPointer<uint32_t>(ATTR_CMP_MASK_MODE_INDEX);	282	opParamInfo_.cmpMaskMode = attrs->GetAttrPointer<uint32_t>(ATTR_CMP_MASK_MODE_INDEX);
273	- opParamInfo_.oriKvStride0 = attrs->GetAttrPointer<uint32_t>(ATTR_ORI_KV_STRIDE_INDEX);
274	- opParamInfo_.cmpKvStride0 = attrs->GetAttrPointer<uint32_t>(ATTR_CMP_KV_STRIDE_INDEX);
275	opParamInfo_.oriWinLeft = attrs->GetAttrPointer<uint32_t>(ATTR_ORI_WIN_LEFT_INDEX);	283	opParamInfo_.oriWinLeft = attrs->GetAttrPointer<uint32_t>(ATTR_ORI_WIN_LEFT_INDEX);
276	opParamInfo_.oriWinRight = attrs->GetAttrPointer<uint32_t>(ATTR_ORI_WIN_RIGHT_INDEX);	284	opParamInfo_.oriWinRight = attrs->GetAttrPointer<uint32_t>(ATTR_ORI_WIN_RIGHT_INDEX);
277	opParamInfo_.layoutQ = attrs->GetStr(ATTR_LAYOUT_Q_INDEX);	285	opParamInfo_.layoutQ = attrs->GetStr(ATTR_LAYOUT_Q_INDEX);
	@@ -293,8 +301,42 @@ ge::graphStatus SMLAInfoParser::GetOpParaInfo()
293	return ge::GRAPH_SUCCESS;	301	return ge::GRAPH_SUCCESS;
294	}	302	}
295		303
296	-ge::graphStatus SMLAInfoParser::GetInOutDataType()	304	+uint64_t SMLAInfoParser::GetOptionalInputStride0(uint32_t inputIndex) const
297	-{	305	+{
		306	+ const gert::Tensor *inputTensor = nullptr;
		307	+ if (inputIndex == ORI_KV_INDEX) {
		308	+ inputTensor = opParamInfo_.oriKv.tensor;
		309	+ } else if (inputIndex == CMP_KV_INDEX) {
		310	+ inputTensor = opParamInfo_.cmpKv.tensor;
		311	+ }
		312	+ if (inputTensor == nullptr) {
		313	+ return 0ULL;
		314	+ }
		315	+
		316	+ // Keep the tiling-context shape path here for future CANN packages that expose real input stride.
		S S songjionghui6月18日【一般】无用代码删除
		317	+ // const gert::StorageShape *inputShape = context_->GetOptionalInputShape(inputIndex);
		318	+ // if (inputShape == nullptr) {
		319	+ // return 0ULL;
		320	+ // }
		321	+ // const gert::Shape &storageShape = inputShape->GetStorageShape();
		322	+
		323	+ const gert::Shape &storageShape = inputTensor->GetStorageShape();
		324	+ if (storageShape.GetDimNum() <= DIM_NUM_ONE) {
		325	+ return 0ULL;
		326	+ }
		327	+
		328	+ uint64_t stride0 = 1ULL;
		329	+ for (size_t i = 1; i < storageShape.GetDimNum(); ++i) {
		330	+ int64_t dim = storageShape.GetDim(i);
		331	+ if (dim <= 0) {
		332	+ return 0ULL;
		333	+ }
		334	+ stride0 *= static_cast<uint64_t>(dim);
		335	+ }
		336	+ return stride0;
		337	+}
		338	+ge::graphStatus SMLAInfoParser::GetInOutDataType()
		339	+{
298	qType_ = opParamInfo_.q.desc->GetDataType();	340	qType_ = opParamInfo_.q.desc->GetDataType();
299	outputType_ = opParamInfo_.attnOut.desc->GetDataType();	341	outputType_ = opParamInfo_.attnOut.desc->GetDataType();
300	if (opParamInfo_.oriKv.desc != nullptr) {	342	if (opParamInfo_.oriKv.desc != nullptr) {
	@@ -361,7 +403,7 @@ ge::graphStatus SMLAInfoParser::GetQueryAndOutLayout()
361	ge::graphStatus SMLAInfoParser::GetKvLayout()	403	ge::graphStatus SMLAInfoParser::GetKvLayout()
362	{	404	{
363	const map<string, SMLALayout> layoutKVMap = {	405	const map<string, SMLALayout> layoutKVMap = {
364	- {"PA_BBND", SMLALayout::PA_BBND},	406	+ {"PA_BBND", SMLALayout::PA_BBND},
365	{"TND", SMLALayout::TND},	407	{"TND", SMLALayout::TND},
366	{"BSND", SMLALayout::BSND},	408	{"BSND", SMLALayout::BSND},
367	};	409	};
	@@ -554,19 +596,13 @@ ge::graphStatus SMLAInfoParser::GetMaxBlockNumPerBatch()
554	OP_LOGE(opName_, "the dim num of cmp_block_table is %u, it should be %u.", cmpDimNum, DIM_NUM_TWO);	596	OP_LOGE(opName_, "the dim num of cmp_block_table is %u, it should be %u.", cmpDimNum, DIM_NUM_TWO);
555	return ge::GRAPH_FAILED;	597	return ge::GRAPH_FAILED;
556	}	598	}
557	- if (qLayout_ == SMLALayout::TND) {	599	+ if (qLayout_ == SMLALayout::TND \|\| qLayout_ == SMLALayout::BSND) {
558	- if (opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(0) != bSize_ - 1) {	600	+ if (opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(0) != bSize_) {
559	- OP_LOGE(opName_, "cmp_block_table's first dimension(%u) should be equal to query's B(%u).",	601	+ OP_LOGE(opName_, "cmp_block_table's first dimension(%lld) should be equal to query's B(%u).",
560	- opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(1), bSize_ - 1);	602	+ opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(0), bSize_);
561	- return ge::GRAPH_FAILED;	603	+ return ge::GRAPH_FAILED;
562	- }	604	+ }
563	- } else if (qLayout_ == SMLALayout::BSND) {	605	+ }
564	- if (opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(0) != bSize_) {
565	- OP_LOGE(opName_, "cmp_block_table's first dimension(%u) should be equal to query's B(%u).",
566	- opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(1), bSize_);
567	- return ge::GRAPH_FAILED;
568	- }
569	- }
570	if (opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(1) <= 0) {	606	if (opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(1) <= 0) {
571	OP_LOGE(opName_, "%s's second dimension(%lld) should be greater than 0",	607	OP_LOGE(opName_, "%s's second dimension(%lld) should be greater than 0",
572	CMP_BLOCK_TABLE_NAME.c_str(), opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(1));	608	CMP_BLOCK_TABLE_NAME.c_str(), opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(1));
	@@ -604,7 +640,7 @@ ge::graphStatus SMLAInfoParser::GetS2Size()
604	s2Size_ = GetAxisNum(oriKvShape_, SMLAAxis::S, kvLayout_);	640	s2Size_ = GetAxisNum(oriKvShape_, SMLAAxis::S, kvLayout_);
605	cmpS2Size_ = GetAxisNum(cmpKvShape_, SMLAAxis::S, kvLayout_);	641	cmpS2Size_ = GetAxisNum(cmpKvShape_, SMLAAxis::S, kvLayout_);
606	return ge::GRAPH_SUCCESS;	642	return ge::GRAPH_SUCCESS;
607	- } else if (kvLayout_ == SMLALayout::PA_BBND) {	643	+ } else if (kvLayout_ == SMLALayout::PA_BBND) {
608	// 获取S2基准PAGE_ATTENTION S2 = block_table.dim1 * block_size	644	// 获取S2基准PAGE_ATTENTION S2 = block_table.dim1 * block_size
609	return GetS2SizeForPageAttention();	645	return GetS2SizeForPageAttention();
610	}	646	}
	@@ -670,8 +706,8 @@ ge::graphStatus SMLAInfoParser::GetActualseqInfo()
670	actualLenDimsQ_ = opParamInfo_.seqUsedQ.tensor->GetShapeSize();	706	actualLenDimsQ_ = opParamInfo_.seqUsedQ.tensor->GetShapeSize();
671	}	707	}
672	}	708	}
673	- if (kvLayout_ != SMLALayout::PA_BBND && kvLayout_ != SMLALayout::BSND && kvLayout_ != SMLALayout::TND) {	709	+ if (kvLayout_ != SMLALayout::PA_BBND && kvLayout_ != SMLALayout::BSND && kvLayout_ != SMLALayout::TND) {
674	- OP_LOGE(opName_, "ori_kv and cmp_kv only support PA_BBND, BSND and TND layout.");	710	+ OP_LOGE(opName_, "ori_kv and cmp_kv only support PA_BBND, BSND and TND layout.");
675	return ge::GRAPH_FAILED;	711	return ge::GRAPH_FAILED;
676	}	712	}
677	if (opParamInfo_.sequsedOriKv.tensor != nullptr) {	713	if (opParamInfo_.sequsedOriKv.tensor != nullptr) {
	@@ -693,8 +729,16 @@ ge::graphStatus SMLAInfoParser::GetActualseqInfo()
693	return ge::GRAPH_FAILED;	729	return ge::GRAPH_FAILED;
694	}	730	}
695	}	731	}
		732	+ if (IsNonEmptyOptionalTensor(opParamInfo_.oriTopkLength.tensor)) {
		733	+ OP_LOGE(opName_, "ori_topk_length is reserved and does not support non-empty tensor in current version");
		734	+ return ge::GRAPH_FAILED;
		735	+ }
		736	+ if (IsNonEmptyOptionalTensor(opParamInfo_.cmpTopkLength.tensor)) {
		737	+ OP_LOGE(opName_, "cmp_topk_length is reserved and does not support non-empty tensor in current version");
		738	+ return ge::GRAPH_FAILED;
		739	+ }
696		740
697	- if (kvLayout_ == SMLALayout::PA_BBND) {	741	+ if (kvLayout_ == SMLALayout::PA_BBND) {
698	if (opParamInfo_.sequsedOriKv.tensor != nullptr) {	742	if (opParamInfo_.sequsedOriKv.tensor != nullptr) {
699	if (qLayout_ == SMLALayout::BSND) {	743	if (qLayout_ == SMLALayout::BSND) {
700	if (opParamInfo_.sequsedOriKv.tensor->GetShapeSize() != bSize_) {	744	if (opParamInfo_.sequsedOriKv.tensor->GetShapeSize() != bSize_) {
	@@ -711,17 +755,18 @@ ge::graphStatus SMLAInfoParser::GetActualseqInfo()
711	}	755	}
712	actualLenDimsKV_ = opParamInfo_.sequsedOriKv.tensor->GetShapeSize();	756	actualLenDimsKV_ = opParamInfo_.sequsedOriKv.tensor->GetShapeSize();
713	} else {	757	} else {
714	- OP_LOGE(opName_, "Input sequsedOriKv must be provided when kv layout is PA_BBND");	758	+ OP_LOGE(opName_, "Input sequsedOriKv must be provided when kv layout is PA_BBND");
715	return ge::GRAPH_FAILED;	759	return ge::GRAPH_FAILED;
716	}	760	}
717	} else if (kvLayout_ == SMLALayout::TND) {	761	} else if (kvLayout_ == SMLALayout::TND) {
718	OP_CHECK_IF(opParamInfo_.sequsedOriKv.tensor != nullptr,	762	OP_CHECK_IF(opParamInfo_.sequsedOriKv.tensor != nullptr,
719	OP_LOGE(opName_, "Input sequsedOriKv must not be provided when kv layout is TND"),	763	OP_LOGE(opName_, "Input sequsedOriKv must not be provided when kv layout is TND"),
720	return ge::GRAPH_FAILED);	764	return ge::GRAPH_FAILED);
721	- } else if (kvLayout_ == SMLALayout::BSND) {	765	+ } else if (kvLayout_ == SMLALayout::BSND) {
722	- } else {	766	+ actualLenDimsKV_ = actualLenDimsOriKV_;
723	- OP_LOGE(opName_, "oriKV and cmpKv only support PA_BBND, TND and BSND layout, but got %d.", kvLayout_);	767	+ } else {
724	- return ge::GRAPH_FAILED;	768	+ OP_LOGE(opName_, "oriKV and cmpKv only support PA_BBND, TND and BSND layout, but got %d.", kvLayout_);
		769	+ return ge::GRAPH_FAILED;
725	}	770	}
726	return ge::GRAPH_SUCCESS;	771	return ge::GRAPH_SUCCESS;
727	}	772	}
	@@ -741,14 +786,15 @@ void SMLAInfoParser::GenerateInfo(SMLATilingInfo &smlaInfo)
741	smlaInfo.cmpS2Size = cmpS2Size_;	786	smlaInfo.cmpS2Size = cmpS2Size_;
742	smlaInfo.gSize = gSize_;	787	smlaInfo.gSize = gSize_;
743	smlaInfo.qHeadDim = qHeadDim_;	788	smlaInfo.qHeadDim = qHeadDim_;
744	- smlaInfo.oriKvHeadDim = oriKvHeadDim_;	789	+ smlaInfo.oriKvHeadDim = oriKvHeadDim_;
745	- smlaInfo.cmpKvHeadDim = cmpKvHeadDim_;	790	+ smlaInfo.cmpKvHeadDim = cmpKvHeadDim_;
746	- smlaInfo.qTSize = qTSize_;	791	+ smlaInfo.qTSize = qTSize_;
747	- smlaInfo.oriSparseBlockCount = oriSparseBlockCount_;	792	+ smlaInfo.oriSparseBlockCount = oriSparseBlockCount_;
748	- smlaInfo.cmpSparseBlockCount = cmpSparseBlockCount_;	793	+ smlaInfo.cmpSparseBlockCount = cmpSparseBlockCount_;
749	- smlaInfo.oriWinLeft = oriWinLeft_;	794	+ smlaInfo.sparseBlockCount = cmpSparseBlockCount_;
750	- smlaInfo.oriWinRight = oriWinRight_;	795	+ smlaInfo.oriWinLeft = oriWinLeft_;
751	- smlaInfo.qType = qType_;	796	+ smlaInfo.oriWinRight = oriWinRight_;
		797	+ smlaInfo.qType = qType_;
752	smlaInfo.oriKvType = oriKvType_;	798	smlaInfo.oriKvType = oriKvType_;
753	smlaInfo.cmpKvType = cmpKvType_;	799	smlaInfo.cmpKvType = cmpKvType_;
754	smlaInfo.outputType = outputType_;	800	smlaInfo.outputType = outputType_;
	@@ -770,9 +816,11 @@ void SMLAInfoParser::GenerateInfo(SMLATilingInfo &smlaInfo)
770	smlaInfo.isSameSeqAllKVTensor = isSameSeqAllKVTensor_;	816	smlaInfo.isSameSeqAllKVTensor = isSameSeqAllKVTensor_;
771		817
772	smlaInfo.softmaxScale = *opParamInfo_.softmaxScale;	818	smlaInfo.softmaxScale = *opParamInfo_.softmaxScale;
773	- smlaInfo.cmpRatio = *opParamInfo_.cmpRatio;	819	+ smlaInfo.cmpRatio = *opParamInfo_.cmpRatio;
774	smlaInfo.oriMaskMode = *opParamInfo_.oriMaskMode;	820	smlaInfo.oriMaskMode = *opParamInfo_.oriMaskMode;
775	smlaInfo.cmpMaskMode = *opParamInfo_.cmpMaskMode;	821	smlaInfo.cmpMaskMode = *opParamInfo_.cmpMaskMode;
		822	+ smlaInfo.oriKvStride0 = GetOptionalInputStride0(ORI_KV_INDEX);
		823	+ smlaInfo.cmpKvStride0 = GetOptionalInputStride0(CMP_KV_INDEX);
776	smlaInfo.oriWinLeft = *opParamInfo_.oriWinLeft;	824	smlaInfo.oriWinLeft = *opParamInfo_.oriWinLeft;
777	smlaInfo.oriWinRight = *opParamInfo_.oriWinRight;	825	smlaInfo.oriWinRight = *opParamInfo_.oriWinRight;
778		826
	@@ -1033,9 +1081,9 @@ ge::graphStatus SMLATilingCheck::CheckSingleParaCuSeqLensOriKv() const
1033	if (opParamInfo_.cuSeqLensOriKv.tensor == nullptr) {	1081	if (opParamInfo_.cuSeqLensOriKv.tensor == nullptr) {
1034	return ge::GRAPH_SUCCESS;	1082	return ge::GRAPH_SUCCESS;
1035	}	1083	}
1036	- if (ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.sequsedOriKv.desc, CU_SEQLENS_ORI_KV_NAME)) {	1084	+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.cuSeqLensOriKv.desc, CU_SEQLENS_ORI_KV_NAME)) {
1037	- return ge::GRAPH_FAILED;	1085	+ return ge::GRAPH_FAILED;
1038	- }	1086	+ }
1039	OP_CHECK_IF(opParamInfo_.cuSeqLensOriKv.tensor->GetShapeSize() != bSize_ + 1,	1087	OP_CHECK_IF(opParamInfo_.cuSeqLensOriKv.tensor->GetShapeSize() != bSize_ + 1,
1040	OP_LOGE(opName_, "Input cuSeqLensOriKv's shapeSize is not equal to B + 1: %u, it is %ld", bSize_ + 1,	1088	OP_LOGE(opName_, "Input cuSeqLensOriKv's shapeSize is not equal to B + 1: %u, it is %ld", bSize_ + 1,
1041	opParamInfo_.cuSeqLensOriKv.tensor->GetShapeSize()),	1089	opParamInfo_.cuSeqLensOriKv.tensor->GetShapeSize()),
	@@ -1048,9 +1096,9 @@ ge::graphStatus SMLATilingCheck::CheckSingleParaCuSeqLensCmpKv() const
1048	if (opParamInfo_.cuSeqLensCmpKv.tensor == nullptr) {	1096	if (opParamInfo_.cuSeqLensCmpKv.tensor == nullptr) {
1049	return ge::GRAPH_SUCCESS;	1097	return ge::GRAPH_SUCCESS;
1050	}	1098	}
1051	- if (ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.sequsedOriKv.desc, CU_SEQLENS_CMP_KV_NAME)) {	1099	+ if (ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.cuSeqLensCmpKv.desc, CU_SEQLENS_CMP_KV_NAME)) {
1052	- return ge::GRAPH_FAILED;	1100	+ return ge::GRAPH_FAILED;
1053	- }	1101	+ }
1054	OP_CHECK_IF(opParamInfo_.cuSeqLensCmpKv.tensor->GetShapeSize() != bSize_ + 1,	1102	OP_CHECK_IF(opParamInfo_.cuSeqLensCmpKv.tensor->GetShapeSize() != bSize_ + 1,
1055	OP_LOGE(opName_, "Input cuSeqLensCmpKv's shapeSize is not equal to B + 1: %u, it is %ld", bSize_ + 1,	1103	OP_LOGE(opName_, "Input cuSeqLensCmpKv's shapeSize is not equal to B + 1: %u, it is %ld", bSize_ + 1,
1056	opParamInfo_.cuSeqLensCmpKv.tensor->GetShapeSize()),	1104	opParamInfo_.cuSeqLensCmpKv.tensor->GetShapeSize()),
	@@ -1096,7 +1144,7 @@ ge::graphStatus SMLATilingCheck::CheckSingleParaCmpSparseIndices() const
1096		1144
1097	ge::graphStatus SMLATilingCheck::CheckSingleParaOriBlockTable() const	1145	ge::graphStatus SMLATilingCheck::CheckSingleParaOriBlockTable() const
1098	{	1146	{
1099	- if (kvLayout_ != SMLALayout::PA_BBND) {	1147	+ if (kvLayout_ != SMLALayout::PA_BBND) {
1100	return ge::GRAPH_SUCCESS;	1148	return ge::GRAPH_SUCCESS;
1101	}	1149	}
1102	const std::vector<size_t> oriBlockTableDimNumList = {DIM_NUM_TWO};	1150	const std::vector<size_t> oriBlockTableDimNumList = {DIM_NUM_TWO};
	@@ -1115,7 +1163,7 @@ ge::graphStatus SMLATilingCheck::CheckSingleParaOriBlockTable() const
1115		1163
1116	ge::graphStatus SMLATilingCheck::CheckSingleParaCmpBlockTable() const	1164	ge::graphStatus SMLATilingCheck::CheckSingleParaCmpBlockTable() const
1117	{	1165	{
1118	- if (kvLayout_ != SMLALayout::PA_BBND) {	1166	+ if (kvLayout_ != SMLALayout::PA_BBND) {
1119	return ge::GRAPH_SUCCESS;	1167	return ge::GRAPH_SUCCESS;
1120	}	1168	}
1121	if (smlaInfo_.perfMode == optiling::SMLATemplateMode::SCFA_TEMPLATE_MODE \|\|	1169	if (smlaInfo_.perfMode == optiling::SMLATemplateMode::SCFA_TEMPLATE_MODE \|\|
	@@ -1163,7 +1211,7 @@ ge::graphStatus SMLATilingCheck::CheckSingleParaMetadata() const
1163	return ge::GRAPH_FAILED;	1211	return ge::GRAPH_FAILED;
1164	}	1212	}
1165	OP_CHECK_IF((opParamInfo_.metadata.tensor->GetShapeSize() != METADATA_LIMIT),	1213	OP_CHECK_IF((opParamInfo_.metadata.tensor->GetShapeSize() != METADATA_LIMIT),
1166	- OP_LOGE(opName_, "input metadata dim 0 must be %u.", METADATA_LIMIT),	1214	+ OP_LOGE(opName_, "input metadata dim 0 must be %u.", METADATA_LIMIT),
1167	return ge::GRAPH_FAILED);	1215	return ge::GRAPH_FAILED);
1168	OP_CHECK_IF(opParamInfo_.metadata.desc->GetDataType() != ge::DT_INT32,	1216	OP_CHECK_IF(opParamInfo_.metadata.desc->GetDataType() != ge::DT_INT32,
1169	OP_LOGE(opName_, "metadata's dtype must be DT_INT32."),	1217	OP_LOGE(opName_, "metadata's dtype must be DT_INT32."),
	@@ -1171,22 +1219,19 @@ ge::graphStatus SMLATilingCheck::CheckSingleParaMetadata() const
1171	return ge::GRAPH_SUCCESS;	1219	return ge::GRAPH_SUCCESS;
1172	}	1220	}
1173		1221
1174	-ge::graphStatus SMLATilingCheck::CheckSingleParaCmpRatio() const	1222	+ge::graphStatus SMLATilingCheck::CheckSingleParaCmpRatio() const
1175	-{	1223	+{
1176	- if (npuArch_ == NpuArch::DAV_3510) {	1224	+ if (npuArch_ == NpuArch::DAV_3510) {
1177	- OP_CHECK_IF(cmpRatio_ < 1 \|\| cmpRatio_ > 128,	1225	+ OP_CHECK_IF(cmpRatio_ < 1 \|\| cmpRatio_ > 128,
1178	- OP_LOGE(opName_, "cmpRatio should be in range [1, 128], but got %u", cmpRatio_),	1226	+ OP_LOGE(opName_, "cmpRatio should be in range [1, 128], but got %ld", cmpRatio_),
1179	- return ge::GRAPH_FAILED);	1227	+ return ge::GRAPH_FAILED);
1180	- } else {	1228	+ } else {
		S S songjionghui6月18日【一般】NPU ARCH上的拦截，建议把具体的版本信息加上，避免后续新增arch版本的时候走入错误的分支
1181	- if (smlaInfo_.perfMode == SMLATemplateMode::CFA_TEMPLATE_MODE \|\|	1229	+ OP_CHECK_IF(cmpRatio_ != 1 && cmpRatio_ != 4 && cmpRatio_ != 128,
1182	- smlaInfo_.perfMode == SMLATemplateMode::SCFA_TEMPLATE_MODE) {	1230	+ OP_LOGE(opName_, "cmpRatio should be 1, 4 or 128, but got %ld", cmpRatio_),
1183	- OP_CHECK_IF(cmpRatio_ != 128 && cmpRatio_ != 4,	1231	+ return ge::GRAPH_FAILED);
1184	- OP_LOGE(opName_, "cmpRatio should be 128 or 4, but got %u", cmpRatio_),	1232	+ }
1185	- return ge::GRAPH_FAILED);	1233	+ return ge::GRAPH_SUCCESS;
1186	- }	1234	+}
1187	- }
1188	- return ge::GRAPH_SUCCESS;
1189	-}
1190		1235
1191	ge::graphStatus SMLATilingCheck::CheckSingleParaOriMaskMode() const	1236	ge::graphStatus SMLATilingCheck::CheckSingleParaOriMaskMode() const
1192	{	1237	{
	@@ -1208,12 +1253,14 @@ ge::graphStatus SMLATilingCheck::CheckSingleParaOriWinRight() const
1208	return ge::GRAPH_SUCCESS;	1253	return ge::GRAPH_SUCCESS;
1209	}	1254	}
1210		1255
1211	-ge::graphStatus SMLATilingCheck::CheckSingleParaCmpResidualKv() const	1256	+ge::graphStatus SMLATilingCheck::CheckSingleParaCmpResidualKv() const
1212	-{	1257	+{
1213	- if (opParamInfo_.cmpMaskMode == 3 && opParamInfo_.cmpRatio != 1) {	1258	+ bool isCmpTemplate = smlaInfo_.perfMode == SMLATemplateMode::CFA_TEMPLATE_MODE \|\|
1214	- OP_CHECK_IF(opParamInfo_.cmpResidualKv.tensor == nullptr,	1259	+ smlaInfo_.perfMode == SMLATemplateMode::SCFA_TEMPLATE_MODE;
1215	- OP_LOGE(opName_, "cmp_redisual_kv is required when cmp_mask_mode=3 and cmp_ratio != 1"),	1260	+ if (isCmpTemplate && *opParamInfo_.cmpMaskMode == 3 && cmpRatio_ != 1) {
1216	- return ge::GRAPH_FAILED);	1261	+ OP_CHECK_IF(opParamInfo_.cmpResidualKv.tensor == nullptr,
		1262	+ OP_LOGE(opName_, "cmp_redisual_kv is required when cmp_mask_mode=3 and cmp_ratio != 1"),
		1263	+ return ge::GRAPH_FAILED);
1217	}	1264	}
1218	return ge::GRAPH_SUCCESS;	1265	return ge::GRAPH_SUCCESS;
1219	}	1266	}
	@@ -1223,19 +1270,19 @@ ge::graphStatus SMLATilingCheck::CheckSinglePara() const
1223	if (	1270	if (
1224	ge::GRAPH_SUCCESS != CheckSingleParaQuery() \|\|	1271	ge::GRAPH_SUCCESS != CheckSingleParaQuery() \|\|
1225	ge::GRAPH_SUCCESS != CheckSingleParaOriKv() \|\|	1272	ge::GRAPH_SUCCESS != CheckSingleParaOriKv() \|\|
1226	- ge::GRAPH_SUCCESS != CheckSingleParaCmpKv() \|\|	1273	+ ge::GRAPH_SUCCESS != CheckSingleParaCmpKv() \|\|
1227	- ge::GRAPH_SUCCESS != CheckSingleParaCuSeqLensOriKv() \|\|	1274	+ ge::GRAPH_SUCCESS != CheckSingleParaCuSeqLensOriKv() \|\|
1228	- ge::GRAPH_SUCCESS != CheckSingleParaCuSeqLensCmpKv() \|\|	1275	+ ge::GRAPH_SUCCESS != CheckSingleParaCuSeqLensCmpKv() \|\|
1229	- ge::GRAPH_SUCCESS != CheckSingleParaCmpResidualKv() \|\|	1276	+ ge::GRAPH_SUCCESS != CheckSingleParaCmpRatio() \|\|
1230	- ge::GRAPH_SUCCESS != CheckSingleParaNumHeads() \|\|	1277	+ ge::GRAPH_SUCCESS != CheckSingleParaCmpResidualKv() \|\|
1231	- ge::GRAPH_SUCCESS != CheckSingleParaKvHeadNums() \|\|	1278	+ ge::GRAPH_SUCCESS != CheckSingleParaNumHeads() \|\|
		1279	+ ge::GRAPH_SUCCESS != CheckSingleParaKvHeadNums() \|\|
1232	ge::GRAPH_SUCCESS != CheckSingleParaCmpSparseIndices() \|\|	1280	ge::GRAPH_SUCCESS != CheckSingleParaCmpSparseIndices() \|\|
1233	ge::GRAPH_SUCCESS != CheckSingleParaOriBlockTable() \|\|	1281	ge::GRAPH_SUCCESS != CheckSingleParaOriBlockTable() \|\|
1234	- ge::GRAPH_SUCCESS != CheckSingleParaCmpBlockTable() \|\|	1282	+ ge::GRAPH_SUCCESS != CheckSingleParaCmpBlockTable() \|\|
1235	- ge::GRAPH_SUCCESS != CheckSingleParaSinks() \|\|	1283	+ ge::GRAPH_SUCCESS != CheckSingleParaSinks() \|\|
1236	- ge::GRAPH_SUCCESS != CheckSingleParaMetadata() \|\|	1284	+ ge::GRAPH_SUCCESS != CheckSingleParaMetadata() \|\|
1237	- ge::GRAPH_SUCCESS != CheckSingleParaCmpRatio() \|\|	1285	+ ge::GRAPH_SUCCESS != CheckSingleParaOriMaskMode() \|\|
1238	- ge::GRAPH_SUCCESS != CheckSingleParaOriMaskMode() \|\|
1239	ge::GRAPH_SUCCESS != CheckSingleParaCmpMaskMode() \|\|	1286	ge::GRAPH_SUCCESS != CheckSingleParaCmpMaskMode() \|\|
1240	ge::GRAPH_SUCCESS != CheckSingleParaOriWinLeft() \|\|	1287	ge::GRAPH_SUCCESS != CheckSingleParaOriWinLeft() \|\|
1241	ge::GRAPH_SUCCESS != CheckSingleParaOriWinRight()) {	1288	ge::GRAPH_SUCCESS != CheckSingleParaOriWinRight()) {
	@@ -1299,7 +1346,7 @@ ge::graphStatus SMLATilingCheck::CheckParaExistence() const
1299	OP_LOGE(opName_, "cuSeqLensOriKv must be provided when kv layout is TND"),	1346	OP_LOGE(opName_, "cuSeqLensOriKv must be provided when kv layout is TND"),
1300	return ge::GRAPH_FAILED);	1347	return ge::GRAPH_FAILED);
1301	} else {	1348	} else {
1302	- if (kvLayout_ == SMLALayout::PA_BBND) {	1349	+ if (kvLayout_ == SMLALayout::PA_BBND) {
1303	std::map<std::string, const void *> ParamExistMap = {	1350	std::map<std::string, const void *> ParamExistMap = {
1304	{"actualSeqLengths", opParamInfo_.sequsedOriKv.tensor},	1351	{"actualSeqLengths", opParamInfo_.sequsedOriKv.tensor},
1305	{"oriBlockTable", opParamInfo_.oriBlockTable.tensor},	1352	{"oriBlockTable", opParamInfo_.oriBlockTable.tensor},
	@@ -1323,31 +1370,28 @@ ge::graphStatus SMLATilingCheck::CheckFeatureShape() const
1323	OP_LOGE(opName_, "T_size of query should be greater than 0, but got %u", qTSize_),	1370	OP_LOGE(opName_, "T_size of query should be greater than 0, but got %u", qTSize_),
1324	return ge::GRAPH_FAILED);	1371	return ge::GRAPH_FAILED);
1325		1372
1326	- if (npuArch_ == NpuArch::DAV_3510) {	1373	+ OP_CHECK_IF(n1Size_ < 1 \|\| n1Size_ > 128,
1327	- OP_CHECK_IF(n1Size_ < 1 \|\| n1Size_ > 128,	1374	+ OP_LOGE(opName_, "q_head_num should be in [1, 128], but got %u", n1Size_),
1328	- OP_LOGE(opName_, "q_head_num should be in [1, 128], but got %u", n1Size_),	1375	+ return ge::GRAPH_FAILED);
1329	- return ge::GRAPH_FAILED);	1376	+
1330	- OP_CHECK_IF(gSize_ < 1 \|\| gSize_ > 128,	1377	+ OP_CHECK_IF(n2Size_ != 1,
1331	- OP_LOGE(opName_, "group num should be in [1, 128] but got %u", gSize_),	1378	+ OP_LOGE(opName_, "kv_head_num should be 1, but got %u", n2Size_),
1332	- return ge::GRAPH_FAILED);
1333	- } else {
1334	- OP_CHECK_IF(n1Size_ % 4 != 0,
1335	- OP_LOGE(opName_, "q_head_num should be multiple of 4, but got %u", n1Size_),
1336	- return ge::GRAPH_FAILED);
1337	- OP_CHECK_IF(gSize_ % 4 != 0,
1338	- OP_LOGE(opName_, "group num should be multiple of 4, but got %u", gSize_),
1339	- return ge::GRAPH_FAILED);
1340	- }
1341	-
1342	- OP_CHECK_IF(n2Size_ != 1,
1343	- OP_LOGE(opName_, "kv_head_num should be 1, but got %u", n2Size_),
1344	return ge::GRAPH_FAILED);	1379	return ge::GRAPH_FAILED);
1345		1380
1346	- OP_CHECK_IF(n1Size_ % n2Size_ != 0,	1381	+ OP_CHECK_IF(n1Size_ % n2Size_ != 0,
1347	- OP_LOGE(opName_, "q_head_num(%u) must be divisible by kv_head_num(%u)", n1Size_, n2Size_),	1382	+ OP_LOGE(opName_, "q_head_num(%u) must be divisible by kv_head_num(%u)", n1Size_, n2Size_),
1348	- return ge::GRAPH_FAILED);	1383	+ return ge::GRAPH_FAILED);
1349	-	1384	+ if (npuArch_ == NpuArch::DAV_3510) {
1350	- OP_CHECK_IF(qHeadDim_ != DIM_LIMIT,	1385	+ OP_CHECK_IF(gSize_ < 1 \|\| gSize_ > 128,
		1386	+ OP_LOGE(opName_, "group num should be in [1, 128], but got %u", gSize_),
		1387	+ return ge::GRAPH_FAILED);
		1388	+ } else {
		1389	+ OP_CHECK_IF(!IsPowerOfTwoInRange(gSize_, 1, 128),
		1390	+ OP_LOGE(opName_, "group num should be power of two in [1, 128], but got %u", gSize_),
		1391	+ return ge::GRAPH_FAILED);
		1392	+ }
		1393	+
		1394	+ OP_CHECK_IF(qHeadDim_ != DIM_LIMIT,
1351	OP_LOGE(opName_, "q_head_dim only support %u, but got %u", DIM_LIMIT, qHeadDim_),	1395	OP_LOGE(opName_, "q_head_dim only support %u, but got %u", DIM_LIMIT, qHeadDim_),
1352	return ge::GRAPH_FAILED);	1396	return ge::GRAPH_FAILED);
1353	OP_CHECK_IF(oriKvHeadDim_ != DIM_LIMIT,	1397	OP_CHECK_IF(oriKvHeadDim_ != DIM_LIMIT,
	@@ -1614,22 +1658,22 @@ ge::graphStatus SparseFlashMlaTiling::DoOpTiling(SMLATilingInfo *tilingInfo)
1614	tilingData_.baseParams.set_sparseBlockSize(tilingInfo->sparseBlockSize);	1658	tilingData_.baseParams.set_sparseBlockSize(tilingInfo->sparseBlockSize);
1615	tilingData_.baseParams.set_returnSoftmaxLse(tilingInfo->returnSoftmaxLse);	1659	tilingData_.baseParams.set_returnSoftmaxLse(tilingInfo->returnSoftmaxLse);
1616		1660
1617	- tilingData_.cmpParams.set_cmpMaxBlockNumPerBatch(tilingInfo->cmpMaxBlockNumPerBatch);	1661	+ tilingData_.cmpParams.set_cmpMaxBlockNumPerBatch(tilingInfo->cmpMaxBlockNumPerBatch);
1618	- tilingData_.cmpParams.set_cmpRatio(tilingInfo->cmpRatio);	1662	+ tilingData_.cmpParams.set_cmpRatio(tilingInfo->cmpRatio);
1619	- tilingData_.cmpParams.set_cmpMaskMode(tilingInfo->cmpMaskMode);	1663	+ tilingData_.cmpParams.set_cmpMaskMode(tilingInfo->cmpMaskMode);
1620	- tilingData_.cmpParams.set_cmpKvStride0(tilingInfo->cmpKvStride0);	1664	+ tilingData_.cmpParams.set_cmpKvStride0(tilingInfo->cmpKvStride0);
1621	-	1665	+ tilingData_.cmpParams.set_cmpKvSeqSize(tilingInfo->cmpS2Size);
1622	- if (tilingInfo->npuArch == NpuArch::DAV_3510) {	1666	+ tilingData_.baseParams.set_actualLenDimsOriKV(tilingInfo->actualLenDimsOriKV);
1623	- tilingData_.baseParams.set_oriSparseBlockCount(tilingInfo->oriSparseBlockCount);	1667	+ tilingData_.baseParams.set_actualLenDimsCmpKV(tilingInfo->actualLenDimsCmpKV);
1624	- tilingData_.baseParams.set_topkValueMode(tilingInfo->topkValueMode);	1668	+ tilingData_.baseParams.set_cmpResidualKVSize(tilingInfo->cmpResidualKVSize);
1625	- tilingData_.cmpParams.set_cmpSparseBlockCount(tilingInfo->cmpSparseBlockCount);	1669	+
1626	- tilingData_.cmpParams.set_cmpKvSeqSize(tilingInfo->cmpS2Size);	1670	+ if (tilingInfo->npuArch == NpuArch::DAV_3510) {
1627	- tilingData_.baseParams.set_actualLenDimsOriKV(tilingInfo->actualLenDimsOriKV);	1671	+ tilingData_.baseParams.set_oriSparseBlockCount(tilingInfo->oriSparseBlockCount);
1628	- tilingData_.baseParams.set_actualLenDimsCmpKV(tilingInfo->actualLenDimsCmpKV);	1672	+ tilingData_.baseParams.set_topkValueMode(tilingInfo->topkValueMode);
1629	- tilingData_.baseParams.set_cmpResidualKVSize(tilingInfo->cmpResidualKVSize);	1673	+ tilingData_.cmpParams.set_cmpSparseBlockCount(tilingInfo->cmpSparseBlockCount);
1630	- } else {	1674	+ } else {
1631	- tilingData_.cmpParams.set_sparseBlockCount(tilingInfo->sparseBlockCount);	1675	+ tilingData_.cmpParams.set_sparseBlockCount(tilingInfo->sparseBlockCount);
1632	- }	1676	+ }
1633		1677
1634	usedCoreNum_ = aicNum;	1678	usedCoreNum_ = aicNum;
1635	tilingData_.baseParams.set_usedCoreNum(usedCoreNum_);	1679	tilingData_.baseParams.set_usedCoreNum(usedCoreNum_);
	@@ -1639,13 +1683,17 @@ ge::graphStatus SparseFlashMlaTiling::DoOpTiling(SMLATilingInfo *tilingInfo)
1639	uint32_t qLayout = static_cast<uint32_t>(tilingInfo->qLayout);	1683	uint32_t qLayout = static_cast<uint32_t>(tilingInfo->qLayout);
1640	uint32_t inputKvLayout = static_cast<uint32_t>(tilingInfo->kvLayout);	1684	uint32_t inputKvLayout = static_cast<uint32_t>(tilingInfo->kvLayout);
1641		1685
1642	- uint32_t tilingKey;	1686	+ uint32_t tilingKey;
1643	- if (tilingInfo->npuArch == NpuArch::DAV_3510) {	1687	+ uint32_t splitG = 0U;
1644	- tilingKey = GET_TPL_TILING_KEY(0U, qLayout, inputKvLayout, static_cast<uint32_t>(tilingInfo->perfMode),	1688	+ uint32_t headRatioOne = static_cast<uint32_t>(
1645	- static_cast<uint32_t>(tilingInfo->gSize > 64));	1689	+ tilingInfo->npuArch == NpuArch::DAV_2201 &&
1646	- } else {	1690	+ tilingInfo->perfMode == SMLATemplateMode::SCFA_TEMPLATE_MODE &&
1647	- tilingKey = GET_TPL_TILING_KEY(0U, qLayout, inputKvLayout, static_cast<uint32_t>(tilingInfo->perfMode));	1691	+ tilingInfo->gSize == 1U);
1648	- }	1692	+ if (tilingInfo->npuArch == NpuArch::DAV_3510) {
		1693	+ splitG = static_cast<uint32_t>(tilingInfo->gSize > 64);
		1694	+ }
		1695	+ tilingKey = GET_TPL_TILING_KEY(0U, qLayout, inputKvLayout, static_cast<uint32_t>(tilingInfo->perfMode), splitG,
		1696	+ headRatioOne);
1649	context_->SetScheduleMode(1);	1697	context_->SetScheduleMode(1);
1650	context_->SetTilingKey(tilingKey);	1698	context_->SetTilingKey(tilingKey);
1651		1699

Mattention/sparse_flash_mla/op_host/sparse_flash_mla_tiling.h+3-5

	@@ -101,9 +101,7 @@ constexpr uint32_t ATTR_ORI_WIN_RIGHT_INDEX = 5;
101	constexpr uint32_t ATTR_LAYOUT_Q_INDEX = 6;	101	constexpr uint32_t ATTR_LAYOUT_Q_INDEX = 6;
102	constexpr uint32_t ATTR_LAYOUT_KV_INDEX = 7;	102	constexpr uint32_t ATTR_LAYOUT_KV_INDEX = 7;
103	constexpr uint32_t ATTR_TOPK_VALUE_MODE_INDEX = 8; // A2/A3	103	constexpr uint32_t ATTR_TOPK_VALUE_MODE_INDEX = 8; // A2/A3
104	-constexpr uint32_t ATTR_ORI_KV_STRIDE_INDEX = 9; // A2/A3	104	+constexpr uint32_t ATTR_RETURN_SOFTMAX_LSE_INDEX = 9;
105	-constexpr uint32_t ATTR_CMP_KV_STRIDE_INDEX = 10; // A2/A3
106	-constexpr uint32_t ATTR_RETURN_SOFTMAX_LSE_INDEX = 11;
107		105
108	// Dim Index	106	// Dim Index
109	constexpr uint32_t DIM_IDX_ONE = 1;	107	constexpr uint32_t DIM_IDX_ONE = 1;
	@@ -168,6 +166,7 @@ TILING_DATA_FIELD_DEF(uint32_t, s2BaseSize)
168	TILING_DATA_FIELD_DEF(uint32_t, actualLenDimsOriKV)	166	TILING_DATA_FIELD_DEF(uint32_t, actualLenDimsOriKV)
169	TILING_DATA_FIELD_DEF(uint32_t, actualLenDimsCmpKV)	167	TILING_DATA_FIELD_DEF(uint32_t, actualLenDimsCmpKV)
170	TILING_DATA_FIELD_DEF(uint32_t, cmpResidualKVSize)	168	TILING_DATA_FIELD_DEF(uint32_t, cmpResidualKVSize)
		169	+TILING_DATA_FIELD_DEF(uint32_t, kvHeadNum)
171	END_TILING_DATA_DEF	170	END_TILING_DATA_DEF
172	REGISTER_TILING_DATA_CLASS(SparseFlashMlaSwaParamsOp, SparseFlashMlaSwaParams)	171	REGISTER_TILING_DATA_CLASS(SparseFlashMlaSwaParamsOp, SparseFlashMlaSwaParams)
173		172
	@@ -214,8 +213,6 @@ struct SMLAParaInfo {
214	const uint32_t *cmpRatio = nullptr;	213	const uint32_t *cmpRatio = nullptr;
215	const uint32_t *oriMaskMode = nullptr;	214	const uint32_t *oriMaskMode = nullptr;
216	const uint32_t *cmpMaskMode = nullptr;	215	const uint32_t *cmpMaskMode = nullptr;
217	- const uint32_t *oriKvStride0 = nullptr; // A2/A3
218	- const uint32_t *cmpKvStride0 = nullptr; // A2/A3
219	const uint32_t *oriWinLeft = nullptr;	216	const uint32_t *oriWinLeft = nullptr;
220	const uint32_t *oriWinRight = nullptr;	217	const uint32_t *oriWinRight = nullptr;
221	const char *layoutQ = nullptr;	218	const char *layoutQ = nullptr;
	@@ -490,6 +487,7 @@ public:
490	ge::graphStatus GetSparseBlockCount();	487	ge::graphStatus GetSparseBlockCount();
491	ge::graphStatus GetActualseqInfo();	488	ge::graphStatus GetActualseqInfo();
492	ge::graphStatus GetSinks();	489	ge::graphStatus GetSinks();
		490	+ uint64_t GetOptionalInputStride0(uint32_t inputIndex) const;
493	void GenerateInfo(SMLATilingInfo &smlaInfo);	491	void GenerateInfo(SMLATilingInfo &smlaInfo);
494	ge::graphStatus Parse(SMLATilingInfo &smlaInfo);	492	ge::graphStatus Parse(SMLATilingInfo &smlaInfo);
495		493

Mattention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_common.h+7-1

	@@ -39,7 +39,7 @@ enum class SMLA_LAYOUT {
39		39
40	template <typename Q_T, typename KV_T, typename OUT_T, const bool FLASH_DECODE = false,	40	template <typename Q_T, typename KV_T, typename OUT_T, const bool FLASH_DECODE = false,
41	SMLA_LAYOUT LAYOUT_T = SMLA_LAYOUT::BSND, SMLA_LAYOUT KV_LAYOUT_T = SMLA_LAYOUT::PA_BBND,	41	SMLA_LAYOUT LAYOUT_T = SMLA_LAYOUT::BSND, SMLA_LAYOUT KV_LAYOUT_T = SMLA_LAYOUT::PA_BBND,
42	- int TEMPLATE_MODE = 0, typename... Args>	42	+ int TEMPLATE_MODE = 0, const bool HEAD_RATIO_ONE = false, typename... Args>
43	struct SMLAType {	43	struct SMLAType {
44	using queryType = Q_T;	44	using queryType = Q_T;
45	using kvType = KV_T;	45	using kvType = KV_T;
	@@ -49,6 +49,7 @@ struct SMLAType {
49	static constexpr SMLA_LAYOUT kvLayout = KV_LAYOUT_T;	49	static constexpr SMLA_LAYOUT kvLayout = KV_LAYOUT_T;
50	static constexpr bool pageAttention = (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND);	50	static constexpr bool pageAttention = (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND);
51	static constexpr int templateMode = TEMPLATE_MODE;	51	static constexpr int templateMode = TEMPLATE_MODE;
		52	+ static constexpr bool headRatioOne = HEAD_RATIO_ONE;
52	};	53	};
53		54
54	// ================================Util functions==================================	55	// ================================Util functions==================================
	@@ -299,6 +300,9 @@ struct ConstInfo {
299	uint32_t actualLenDimsQ = 0U; // query的actualSeqLength 的维度	300	uint32_t actualLenDimsQ = 0U; // query的actualSeqLength 的维度
300	uint32_t actualLenDimsKV = 0U; // KV 的actualSeqLength 的维度	301	uint32_t actualLenDimsKV = 0U; // KV 的actualSeqLength 的维度
301		302
		303	+ uint32_t actualLenDimsCmpKV = 0U;
		304	+ uint32_t cmpResidualKVSize = 0U;
		305	+
302	// TND	306	// TND
303	uint32_t s2Start = 0U; // TND场景下，S2的起始位置	307	uint32_t s2Start = 0U; // TND场景下，S2的起始位置
304	uint32_t s2End = 0U; // 单核TND场景下S2循环index上限	308	uint32_t s2End = 0U; // 单核TND场景下S2循环index上限
	@@ -321,6 +325,8 @@ struct ConstInfo {
321	// cmp attr	325	// cmp attr
322	int64_t cmpRatio = 0;	326	int64_t cmpRatio = 0;
323		327
		328	+ uint64_t cmpSeqSize = 0ULL;
		329	+
324	// win	330	// win
325	int32_t oriWinRight = 0;	331	int32_t oriWinRight = 0;
326	int32_t oriWinLeft = 128;	332	int32_t oriWinLeft = 128;

Mattention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_metadata.h+3-2

	@@ -26,7 +26,7 @@ constexpr uint32_t AIV_CORE_NUM = 72;
26	constexpr uint32_t SMLA_META_SIZE = 1024;	26	constexpr uint32_t SMLA_META_SIZE = 1024;
27	using SMLA_METADATA_T = int32_t;	27	using SMLA_METADATA_T = int32_t;
28		28
29	-constexpr uint32_t FA_METADATA_SIZE = 8;	29	+constexpr uint32_t FA_METADATA_SIZE = 9;
30	constexpr uint32_t FD_METADATA_SIZE = 8;	30	constexpr uint32_t FD_METADATA_SIZE = 8;
31		31
32	// FA Metadata Index Definitions	32	// FA Metadata Index Definitions
	@@ -38,6 +38,7 @@ constexpr uint32_t FA_BN2_END_INDEX = 4;
38	constexpr uint32_t FA_M_END_INDEX = 5;	38	constexpr uint32_t FA_M_END_INDEX = 5;
39	constexpr uint32_t FA_S2_END_INDEX = 6;	39	constexpr uint32_t FA_S2_END_INDEX = 6;
40	constexpr uint32_t FA_FIRST_FD_DATA_WORKSPACE_IDX_INDEX = 7;	40	constexpr uint32_t FA_FIRST_FD_DATA_WORKSPACE_IDX_INDEX = 7;
		41	+constexpr uint32_t FA_S2_MAX_NUM = 8;
41		42
42	// FD Metadata Index Definitions	43	// FD Metadata Index Definitions
43	constexpr uint32_t FD_CORE_ENABLE_INDEX = 0;	44	constexpr uint32_t FD_CORE_ENABLE_INDEX = 0;
	@@ -76,4 +77,4 @@ struct SasMetadata {
76	static_assert(SMLA_META_SIZE * sizeof(SMLA_METADATA_T) >= sizeof(detail::SasMetadata));	77	static_assert(SMLA_META_SIZE * sizeof(SMLA_METADATA_T) >= sizeof(detail::SasMetadata));
77	} // namespace optiling	78	} // namespace optiling
78		79
79	-#endif	80	+#endif

Mattention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_scfa_block_cube.h+5-3

	@@ -410,7 +410,8 @@ __aicore__ inline void SMLACubeBlock<SMLAT>::ComputeMm1(const RunInfo &info, con
410		410
411	uint32_t headStride = constInfo.headDim;	411	uint32_t headStride = constInfo.headDim;
412	uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim;	412	uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim;
413	- uint32_t batchStride = constInfo.kvSeqSize * seqStride;	413	+ uint64_t batchStride = (constInfo.oriKvStride0 == 0) ?
		414	+ static_cast<uint64_t>(constInfo.kvSeqSize) * seqStride : constInfo.oriKvStride0;
414		415
415	uint32_t curS2 = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint;	416	uint32_t curS2 = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint;
416	uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride + (uint64_t)info.n2Idx * headStride + kL1 * D_SPLIT_SIZE;	417	uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride + (uint64_t)info.n2Idx * headStride + kL1 * D_SPLIT_SIZE;
	@@ -671,7 +672,8 @@ __aicore__ inline void SMLACubeBlock<SMLAT>::ComputeMm2(const RunInfo &info, con
671		672
672	uint32_t headStride = constInfo.headDim;	673	uint32_t headStride = constInfo.headDim;
673	uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim;	674	uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim;
674	- uint32_t batchStride = constInfo.kvSeqSize * seqStride;	675	+ uint64_t batchStride = (constInfo.oriKvStride0 == 0) ?
		676	+ static_cast<uint64_t>(constInfo.kvSeqSize) * seqStride : constInfo.oriKvStride0;
675		677
676	uint32_t curS2 = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint;	678	uint32_t curS2 = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint;
677	uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride + (uint64_t)info.n2Idx * headStride + nL1 * N_SPLIT_SIZE;	679	uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride + (uint64_t)info.n2Idx * headStride + nL1 * N_SPLIT_SIZE;
	@@ -850,4 +852,4 @@ __aicore__ inline void SMLACubeBlock<SMLAT>::ComputeMm2(const RunInfo &info, con
850	qpL1BufIter += mL1Loops;	852	qpL1BufIter += mL1Loops;
851	}	853	}
852	} // namespace SMLAKernel	854	} // namespace SMLAKernel
853	-#endif // SPARSE_FLASH_MLA_SCFA_BLOCK_CUBE_H	855	+#endif // SPARSE_FLASH_MLA_SCFA_BLOCK_CUBE_H

Mattention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_scfa_block_vector.h+134-80

	@@ -120,6 +120,7 @@ private:
120	static constexpr bool FLASH_DECODE = SMLAT::flashDecode;	120	static constexpr bool FLASH_DECODE = SMLAT::flashDecode;
121	static constexpr SMLA_LAYOUT LAYOUT_T = SMLAT::layout;	121	static constexpr SMLA_LAYOUT LAYOUT_T = SMLAT::layout;
122	static constexpr SMLA_LAYOUT KV_LAYOUT_T = SMLAT::kvLayout;	122	static constexpr SMLA_LAYOUT KV_LAYOUT_T = SMLAT::kvLayout;
		123	+ static constexpr bool HEAD_RATIO_ONE = SMLAT::headRatioOne;
123		124
124	static constexpr uint64_t MERGE_CACHE_GM_BUF_NUM = 3;	125	static constexpr uint64_t MERGE_CACHE_GM_BUF_NUM = 3;
125	static constexpr uint64_t SYNC_INPUT_BUF1_FLAG = 2;	126	static constexpr uint64_t SYNC_INPUT_BUF1_FLAG = 2;
	@@ -504,7 +505,9 @@ __aicore__ inline void SMLAVectorBlock<SMLAT>::ProcessVec1SingleBuf(const RunInf
504	uint32_t loopCount = (mSplitInfo.vecDealM + mSplitSize - 1) / mSplitSize;	505	uint32_t loopCount = (mSplitInfo.vecDealM + mSplitSize - 1) / mSplitSize;
505	uint32_t tailSplitSize = mSplitInfo.vecDealM - (loopCount - 1) * mSplitSize;	506	uint32_t tailSplitSize = mSplitInfo.vecDealM - (loopCount - 1) * mSplitSize;
506		507
507	- SliceAndContactSinksValue((mSplitInfo.nBufferStartM + mSplitInfo.vecStartM) % constInfo.qHeadNum,	508	+ uint32_t sinkHeadIdx = (info.n2IdxReal * constInfo.gSize + mSplitInfo.nBufferStartM + mSplitInfo.vecStartM) %
		509	+ constInfo.qHeadNum;
		510	+ SliceAndContactSinksValue(sinkHeadIdx,
508	mSplitInfo.vecDealM);	511	mSplitInfo.vecDealM);
509		512
510	for (uint32_t i = 0, dealSize = mSplitSize; i < loopCount; i++) {	513	for (uint32_t i = 0, dealSize = mSplitSize; i < loopCount; i++) {
	@@ -545,10 +548,17 @@ __aicore__ inline int64_t SMLAVectorBlock<SMLAT>::GetKeyGmOffset(int64_t realS2I
545	blkTableOffset * static_cast<int64_t>(constInfo.kvHeadNum) * static_cast<int64_t>(constInfo.headDim);	548	blkTableOffset * static_cast<int64_t>(constInfo.kvHeadNum) * static_cast<int64_t>(constInfo.headDim);
546		549
547	} else if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::BSND) {	550	} else if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::BSND) {
548	- realKeyGmOffset = runInfo.bIdx * constInfo.kvSeqSize / constInfo.cmpRatio * constInfo.kvHeadNum + realS2Idx * constInfo.kvHeadNum;	551	+ int64_t batchStride = (constInfo.cmpKvStride0 == 0) ?
		552	+ static_cast<int64_t>(constInfo.cmpSeqSize) * static_cast<int64_t>(constInfo.kvHeadNum) *
		553	+ static_cast<int64_t>(constInfo.headDim) :
		554	+ static_cast<int64_t>(constInfo.cmpKvStride0);
		555	+ realKeyGmOffset = static_cast<int64_t>(runInfo.bIdx) * batchStride +
		556	+ realS2Idx * static_cast<int64_t>(constInfo.kvHeadNum) * static_cast<int64_t>(constInfo.headDim) +
		557	+ static_cast<int64_t>(runInfo.n2Idx) * static_cast<int64_t>(constInfo.headDim);
549	} else if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::TND) {	558	} else if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::TND) {
550	- realKeyGmOffset = (runInfo.tensorCmpBOffset + realS2Idx * constInfo.kvHeadNum * constInfo.headDim) /	559	+ realKeyGmOffset = runInfo.tensorCmpBOffset +
551	- constInfo.headDim;	560	+ realS2Idx * static_cast<int64_t>(constInfo.kvHeadNum) * static_cast<int64_t>(constInfo.headDim) +
		561	+ static_cast<int64_t>(runInfo.n2Idx) * static_cast<int64_t>(constInfo.headDim);
552	}	562	}
553	return realKeyGmOffset;	563	return realKeyGmOffset;
554	}	564	}
	@@ -568,16 +578,10 @@ __aicore__ inline void SMLAVectorBlock<SMLAT>::CopyInSingleKv(int64_t &mte2Size,
568	intriParams.blockCount = 1;	578	intriParams.blockCount = 1;
569	intriParams.dstStride = 0;	579	intriParams.dstStride = 0;
570	intriParams.srcStride = 0;	580	intriParams.srcStride = 0;
571	- DataCopyPadExtParams<KV_T> padParams;	581	+ DataCopyPadExtParams<KV_T> padParams{false, 0, 0, 0};
572	- if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND) {	582	+ DataCopyPad(
573	- DataCopyPad(	583	+ kvMergUb_[mergeMte3Idx % 2 * INPUT2_BUFFER_OFFSET / sizeof(KV_T) + (mte2Size - mte3Size) * constInfo.headDim],
574	- kvMergUb_[mergeMte3Idx % 2 * INPUT2_BUFFER_OFFSET / sizeof(KV_T) + (mte2Size - mte3Size) * constInfo.headDim],	584	+ cmpKvGm_[keyBNBOffset], intriParams, padParams);
575	- cmpKvGm_[keyBNBOffset], intriParams, padParams);
576	- } else {
577	- DataCopyPad(
578	- kvMergUb_[mergeMte3Idx % 2 * INPUT2_BUFFER_OFFSET / sizeof(KV_T) + (mte2Size - mte3Size) * constInfo.headDim],
579	- cmpKvGm_[keyBNBOffset * constInfo.headDim], intriParams, padParams);
580	- }
581	mte2Size += validS2Count;	585	mte2Size += validS2Count;
582	}	586	}
583		587
	@@ -594,18 +598,8 @@ __aicore__ inline void SMLAVectorBlock<SMLAT>::CopyInKv(int64_t &mte2Size, int64
594	}	598	}
595		599
596	int64_t keySrcStride = 0;	600	int64_t keySrcStride = 0;
597	- if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND) {	601	+ keySrcStride = ((keyOffset1 > keyOffset2 ? (keyOffset1 - keyOffset2) :
598	- int64_t blkTableSrcStride =	602	+ (keyOffset2 - keyOffset1)) - constInfo.sparseBlockSize * constInfo.headDim) * sizeof(KV_T);
599	- ((keyOffset1 > keyOffset2 ? (keyOffset1 - keyOffset2) :
600	- (keyOffset2 - keyOffset1)) - constInfo.sparseBlockSize * constInfo.headDim);
601	- keySrcStride = blkTableSrcStride * sizeof(KV_T);
602	- } else if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::BSND) {
603	- keySrcStride = ((keyOffset1 > keyOffset2 ? (keyOffset1 - keyOffset2) :
604	- (keyOffset2 - keyOffset1)) - constInfo.sparseBlockSize) * constInfo.headDim * sizeof(KV_T);
605	- } else if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::TND) {
606	- keySrcStride = ((keyOffset1 > keyOffset2 ? (keyOffset1 - keyOffset2) :
607	- (keyOffset2 - keyOffset1)) - constInfo.sparseBlockSize) * constInfo.headDim * sizeof(KV_T);
608	- }
609	if (unlikely(keySrcStride >= INT32_MAX \|\| keySrcStride < 0 \|\|	603	if (unlikely(keySrcStride >= INT32_MAX \|\| keySrcStride < 0 \|\|
610	realS2Idx1 + constInfo.sparseBlockSize >= s2IdLimit \|\|	604	realS2Idx1 + constInfo.sparseBlockSize >= s2IdLimit \|\|
611	realS2Idx2 + constInfo.sparseBlockSize >= s2IdLimit)) {	605	realS2Idx2 + constInfo.sparseBlockSize >= s2IdLimit)) {
	@@ -619,21 +613,15 @@ __aicore__ inline void SMLAVectorBlock<SMLAT>::CopyInKv(int64_t &mte2Size, int64
619	intriParams.blockCount = (keyOffset1 >= 0) + (keyOffset2 >= 0);	613	intriParams.blockCount = (keyOffset1 >= 0) + (keyOffset2 >= 0);
620	intriParams.dstStride = 0;	614	intriParams.dstStride = 0;
621	intriParams.srcStride = keySrcStride;	615	intriParams.srcStride = keySrcStride;
622	- DataCopyPadExtParams<KV_T> padParams;	616	+ DataCopyPadExtParams<KV_T> padParams{false, 0, 0, 0};
623		617
624	int64_t startGmOffset = keyOffset1 > -1 ? keyOffset1 : keyOffset2;	618	int64_t startGmOffset = keyOffset1 > -1 ? keyOffset1 : keyOffset2;
625	if (keyOffset2 > -1 && keyOffset2 < keyOffset1) {	619	if (keyOffset2 > -1 && keyOffset2 < keyOffset1) {
626	startGmOffset = keyOffset2;	620	startGmOffset = keyOffset2;
627	}	621	}
628	- if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND) {	622	+ DataCopyPad(kvMergUb_[mergeMte3Idx % 2 * INPUT2_BUFFER_OFFSET / sizeof(KV_T) +
629	- DataCopyPad(kvMergUb_[mergeMte3Idx % 2 * INPUT2_BUFFER_OFFSET / sizeof(KV_T) +	623	+ (mte2Size - mte3Size) * constInfo.headDim],
630	- (mte2Size - mte3Size) * constInfo.headDim],	624	+ cmpKvGm_[startGmOffset], intriParams, padParams);
631	- cmpKvGm_[startGmOffset], intriParams, padParams);
632	- } else {
633	- DataCopyPad(kvMergUb_[mergeMte3Idx % 2 * INPUT2_BUFFER_OFFSET / sizeof(KV_T) +
634	- (mte2Size - mte3Size) * constInfo.headDim],
635	- cmpKvGm_[startGmOffset * constInfo.headDim], intriParams, padParams);
636	- }
637	mte2Size += ((keyOffset1 > -1) + (keyOffset2 > -1)) * constInfo.sparseBlockSize;	625	mte2Size += ((keyOffset1 > -1) + (keyOffset2 > -1)) * constInfo.sparseBlockSize;
638	}	626	}
639	}	627	}
	@@ -725,7 +713,11 @@ __aicore__ inline void SMLAVectorBlock<SMLAT>::ProcessVec1L(const RunInfo &info)
725	mSplitInfo.vecDealM = mSplitInfo.nBufferDealM - mSplitInfo.vecDealM;	713	mSplitInfo.vecDealM = mSplitInfo.nBufferDealM - mSplitInfo.vecDealM;
726	}	714	}
727		715
728	- CrossCoreWaitFlag(constInfo.syncC1V1);	716	+ if constexpr (HEAD_RATIO_ONE) {
		717	+ CrossCoreWaitFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_FIX>(constInfo.syncC1V1);
		718	+ } else {
		719	+ CrossCoreWaitFlag(constInfo.syncC1V1);
		720	+ }
729	// vec1 compute	721	// vec1 compute
730	ProcessVec1SingleBuf(info, mSplitInfo);	722	ProcessVec1SingleBuf(info, mSplitInfo);
731	CrossCoreSetFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_MTE3>(constInfo.syncV1C2);	723	CrossCoreSetFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_MTE3>(constInfo.syncV1C2);
	@@ -771,7 +763,11 @@ __aicore__ inline void SMLAVectorBlock<SMLAT>::ProcessVec2L(const RunInfo &info)
771	mSplitInfo.vecStartM = mSplitInfo.vecDealM;	763	mSplitInfo.vecStartM = mSplitInfo.vecDealM;
772	mSplitInfo.vecDealM = mSplitInfo.nBufferDealM - mSplitInfo.vecDealM;	764	mSplitInfo.vecDealM = mSplitInfo.nBufferDealM - mSplitInfo.vecDealM;
773	}	765	}
774	- CrossCoreWaitFlag(constInfo.syncC2V2);	766	+ if constexpr (HEAD_RATIO_ONE) {
		767	+ CrossCoreWaitFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_FIX>(constInfo.syncC2V2);
		768	+ } else {
		769	+ CrossCoreWaitFlag(constInfo.syncC2V2);
		770	+ }
775	ProcessVec2SingleBuf(info, mSplitInfo);	771	ProcessVec2SingleBuf(info, mSplitInfo);
776	}	772	}
777	}	773	}
	@@ -899,56 +895,114 @@ __aicore__ inline void SMLAVectorBlock<SMLAT>::DealBmm2ResBaseBlock(const RunInf
899	uint32_t inOutBaseOffset = mStart * columnCount;	895	uint32_t inOutBaseOffset = mStart * columnCount;
900	uint32_t baseOffset = mSplitInfo.nBufferStartM / 2 + startRow;	896	uint32_t baseOffset = mSplitInfo.nBufferStartM / 2 + startRow;
901		897
902	- // 除第一个循环外，均需要更新中间计算结果	898	+ if constexpr (!HEAD_RATIO_ONE) {
903	- if (!info.isFirstSInnerLoop) {	899	+ if (!info.isFirstSInnerLoop) {
904	- event_t eventIdMte2WaitMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2));	900	+ event_t eventIdMte2WaitMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2));
905	- SetFlag<HardEvent::MTE3_MTE2>(eventIdMte2WaitMte3);	901	+ SetFlag<HardEvent::MTE3_MTE2>(eventIdMte2WaitMte3);
906	- WaitFlag<HardEvent::MTE3_MTE2>(eventIdMte2WaitMte3);	902	+ WaitFlag<HardEvent::MTE3_MTE2>(eventIdMte2WaitMte3);
907		903
908	- LocalTensor<MM2_OUT_T> bmm2ResPreUb = inputBuff1.Get<MM2_OUT_T>();	904	+ LocalTensor<MM2_OUT_T> bmm2ResPreUb = inputBuff1.Get<MM2_OUT_T>();
909	- bmm2ResPreUb = bmm2ResPreUb[pingpongFlag * INPUT1_BUFFER_OFFSET / sizeof(MM2_OUT_T)];	905	+ bmm2ResPreUb = bmm2ResPreUb[pingpongFlag * INPUT1_BUFFER_OFFSET / sizeof(MM2_OUT_T)];
910	- WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF1_FLAG + pingpongFlag);	906	+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF1_FLAG + pingpongFlag);
911		907
912	- uint64_t vec2ResGmOffset = ((info.loop - 1) % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + inOutBaseOffset;	908	+ uint64_t vec2ResGmOffset =
913	- DataCopy(bmm2ResPreUb, vec2ResGm[vec2ResGmOffset], vec2ComputeSize);	909	+ ((info.loop - 1) % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + inOutBaseOffset;
		910	+ DataCopy(bmm2ResPreUb, vec2ResGm[vec2ResGmOffset], vec2ComputeSize);
914		911
915	- SetFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_BUF1_FLAG);	912	+ SetFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_BUF1_FLAG);
916	- WaitFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_BUF1_FLAG);	913	+ WaitFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_BUF1_FLAG);
917		914
918	- uint32_t idx = info.loop % (constInfo.preLoadNum);	915	+ uint32_t idx = info.loop % (constInfo.preLoadNum);
919	- LocalTensor<T> expUb = v0ValidSizeBuff.Get<T>()[384]; // sumUb用临时内存 16 * 32B = 512B	916	+ LocalTensor<T> expUb = v0ValidSizeBuff.Get<T>()[384];
920	- Brcb(expUb, softmaxExpUb[idx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset], (dealRowCount + 7) / 8,	917	+ Brcb(expUb, softmaxExpUb[idx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset],
921	- {1, 8});	918	+ (dealRowCount + 7) / 8, {1, 8});
922	- PipeBarrier<PIPE_V>();	919	+ PipeBarrier<PIPE_V>();
923		920
924	- RowMuls(bmm2ResPreUb, bmm2ResPreUb, expUb, dealRowCount, columnCount, actualColumnCount);	921	+ RowMuls(bmm2ResPreUb, bmm2ResPreUb, expUb, dealRowCount, columnCount, actualColumnCount);
925	- AscendC::PipeBarrier<PIPE_V>();	922	+ AscendC::PipeBarrier<PIPE_V>();
926	- Add(bmm2ResUb, bmm2ResUb, bmm2ResPreUb, vec2ComputeSize);	923	+ Add(bmm2ResUb, bmm2ResUb, bmm2ResPreUb, vec2ComputeSize);
927	- AscendC::PipeBarrier<PIPE_V>();	924	+ AscendC::PipeBarrier<PIPE_V>();
928		925
929	- SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF1_FLAG + pingpongFlag);	926	+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF1_FLAG + pingpongFlag);
930	- pingpongFlag ^= 1; // pingpong 0 1 切换	927	+ pingpongFlag ^= 1;
931	- }	928	+ }
932		929
933	- // 最后一次输出计算结果，否则将中间结果暂存至workspace	930	+ if (info.isLastS2Loop) {
934	- if (info.isLastS2Loop) {	931	+ uint32_t idx = info.loop % (constInfo.preLoadNum);
935	- uint32_t idx = info.loop % (constInfo.preLoadNum);	932	+ LocalTensor<T> tmpSumUb = v0ValidSizeBuff.Get<T>()[384];
936	- LocalTensor<T> tmpSumUb = v0ValidSizeBuff.Get<T>()[384]; // sumUb用临时内存 16 * 32B = 512B	933	+ Brcb(tmpSumUb, softmaxSumUb[idx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset],
937	- Brcb(tmpSumUb, softmaxSumUb[idx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset], (dealRowCount + 7) / 8,	934	+ (dealRowCount + 7) / 8, {1, 8});
938	- {1, 8});	935	+ PipeBarrier<PIPE_V>();
939	- PipeBarrier<PIPE_V>();	936	+ RowDivs(bmm2ResUb, bmm2ResUb, tmpSumUb, dealRowCount, columnCount, actualColumnCount);
940	- RowDivs(bmm2ResUb, bmm2ResUb, tmpSumUb, dealRowCount, columnCount, actualColumnCount);	937	+ PipeBarrier<PIPE_V>();
941	- PipeBarrier<PIPE_V>();	938	+ Bmm2ResCopyOut(info, bmm2ResUb, mStart, dealRowCount, columnCount, actualColumnCount);
942	- Bmm2ResCopyOut(info, bmm2ResUb, mStart, dealRowCount, columnCount, actualColumnCount);	939	+ } else {
		940	+ LocalTensor<T> outUb = outputBuff1.Get<T>();
		941	+ WaitFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);
		942	+ DataCopy(outUb, bmm2ResUb, dealRowCount * columnCount);
		943	+ SetFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF1_FLAG);
		944	+ WaitFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF1_FLAG);
		945	+ uint64_t vec2ResGmOffset =
		946	+ (info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + inOutBaseOffset;
		947	+ DataCopy(vec2ResGm[vec2ResGmOffset], outUb, vec2ComputeSize);
		948	+ SetFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);
		949	+ }
943	} else {	950	} else {
944	- LocalTensor<T> outUb = outputBuff1.Get<T>();	951	+ // 除第一个循环外，均需要更新中间计算结果
945	- WaitFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);	952	+ if (!info.isFirstSInnerLoop) {
946	- DataCopy(outUb, bmm2ResUb, dealRowCount * columnCount);	953	+ event_t eventIdMte2WaitMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2));
947	- SetFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF1_FLAG);	954	+ SetFlag<HardEvent::MTE3_MTE2>(eventIdMte2WaitMte3);
948	- WaitFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF1_FLAG);	955	+ WaitFlag<HardEvent::MTE3_MTE2>(eventIdMte2WaitMte3);
949	- uint64_t vec2ResGmOffset = (info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + inOutBaseOffset;	956	+
950	- DataCopy(vec2ResGm[vec2ResGmOffset], outUb, vec2ComputeSize);	957	+ LocalTensor<MM2_OUT_T> bmm2ResPreUb = inputBuff1.Get<MM2_OUT_T>();
951	- SetFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);	958	+ bmm2ResPreUb = bmm2ResPreUb[pingpongFlag * INPUT1_BUFFER_OFFSET / sizeof(MM2_OUT_T)];
		959	+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF1_FLAG + pingpongFlag);
		960	+
		961	+ uint64_t accumGmOffset =
		962	+ ((info.loop - 1) % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + inOutBaseOffset;
		963	+ DataCopy(bmm2ResPreUb, mm2ResGm[accumGmOffset], vec2ComputeSize);
		964	+
		965	+ SetFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_BUF1_FLAG);
		966	+ WaitFlag<AscendC::HardEvent::MTE2_V>(SYNC_INPUT_BUF1_FLAG);
		967	+
		968	+ uint32_t idx = info.loop % (constInfo.preLoadNum);
		969	+ LocalTensor<T> expUb = v0ValidSizeBuff.Get<T>()[384]; // sumUb用临时内存 16 * 32B = 512B
		970	+ Brcb(expUb, softmaxExpUb[idx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset],
		971	+ (dealRowCount + 7) / 8, {1, 8});
		972	+ PipeBarrier<PIPE_V>();
		973	+
		974	+ RowMuls(bmm2ResPreUb, bmm2ResPreUb, expUb, dealRowCount, columnCount, actualColumnCount);
		975	+ AscendC::PipeBarrier<PIPE_V>();
		976	+ Add(bmm2ResUb, bmm2ResUb, bmm2ResPreUb, vec2ComputeSize);
		977	+ AscendC::PipeBarrier<PIPE_V>();
		978	+
		979	+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF1_FLAG + pingpongFlag);
		980	+ pingpongFlag ^= 1; // pingpong 0 1 切换
		981	+ }
		982	+
		983	+ // 最后一次输出计算结果，否则将中间结果暂存至workspace
		984	+ if (info.isLastS2Loop) {
		985	+ uint32_t idx = info.loop % (constInfo.preLoadNum);
		986	+ LocalTensor<T> tmpSumUb = v0ValidSizeBuff.Get<T>()[384]; // sumUb用临时内存 16 * 32B = 512B
		987	+ Brcb(tmpSumUb, softmaxSumUb[idx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset],
		988	+ (dealRowCount + 7) / 8, {1, 8});
		989	+ PipeBarrier<PIPE_V>();
		990	+ RowDivs(bmm2ResUb, bmm2ResUb, tmpSumUb, dealRowCount, columnCount, actualColumnCount);
		991	+ PipeBarrier<PIPE_V>();
		992	+ Bmm2ResCopyOut(info, bmm2ResUb, mStart, dealRowCount, columnCount, actualColumnCount);
		993	+ } else if (!info.isFirstSInnerLoop) {
		994	+ LocalTensor<T> outUb = outputBuff1.Get<T>();
		995	+ WaitFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);
		996	+ DataCopy(outUb, bmm2ResUb, dealRowCount * columnCount);
		997	+ SetFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF1_FLAG);
		998	+ WaitFlag<AscendC::HardEvent::V_MTE3>(SYNC_OUTPUT_BUF1_FLAG);
		999	+ uint64_t accumGmOffset =
		1000	+ (info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + inOutBaseOffset;
		1001	+ DataCopy(mm2ResGm[accumGmOffset], outUb, vec2ComputeSize);
		1002	+ SetFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);
		1003	+ WaitFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);
		1004	+ SetFlag<AscendC::HardEvent::MTE3_V>(SYNC_OUTPUT_BUF1_FLAG);
		1005	+ }
952	}	1006	}
953	}	1007	}
954		1008

Mattention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_scfa_kernel.h+176-56

	@@ -87,8 +87,9 @@ public:
87	__gm__ uint8_t cmpSparseIndices, __gm__ uint8_t oriBlockTable,	87	__gm__ uint8_t cmpSparseIndices, __gm__ uint8_t oriBlockTable,
88	__gm__ uint8_t cmpBlockTable, __gm__ uint8_t cuSeqlensQ,	88	__gm__ uint8_t cmpBlockTable, __gm__ uint8_t cuSeqlensQ,
89	__gm__ uint8_t* cuSeqlensKV, __gm__ uint8_t *cuSeqlensCmpKV,	89	__gm__ uint8_t* cuSeqlensKV, __gm__ uint8_t *cuSeqlensCmpKV,
90	- __gm__ uint8_t seqUsedQ, __gm__ uint8_t seqUsedKV, __gm__ uint8_t *sinks,	90	+ __gm__ uint8_t seqUsedQ, __gm__ uint8_t seqUsedKV,
91	- __gm__ uint8_t metadata, __gm__ uint8_t attentionOut,	91	+ __gm__ uint8_t seqUsedCmpKV, __gm__ uint8_t cmpResidualKV,
		92	+ __gm__ uint8_t sinks, __gm__ uint8_t metadata, __gm__ uint8_t *attentionOut,
92	__gm__ uint8_t softmaxLse, __gm__ uint8_t workspace,	93	__gm__ uint8_t softmaxLse, __gm__ uint8_t workspace,
93	const SparseFlashMlaTilingData __restrict tiling, __gm__ uint8_t gmTiling,	94	const SparseFlashMlaTilingData __restrict tiling, __gm__ uint8_t gmTiling,
94	TPipe *tPipe);	95	TPipe *tPipe);
	@@ -100,6 +101,7 @@ private:
100	static constexpr bool FLASH_DECODE = SMLAT::flashDecode;	101	static constexpr bool FLASH_DECODE = SMLAT::flashDecode;
101	static constexpr SMLA_LAYOUT LAYOUT_T = SMLAT::layout;	102	static constexpr SMLA_LAYOUT LAYOUT_T = SMLAT::layout;
102	static constexpr SMLA_LAYOUT KV_LAYOUT_T = SMLAT::kvLayout;	103	static constexpr SMLA_LAYOUT KV_LAYOUT_T = SMLAT::kvLayout;
		104	+ static constexpr bool HEAD_RATIO_ONE = SMLAT::headRatioOne;
103		105
104	static constexpr uint32_t PRELOAD_NUM = 2;	106	static constexpr uint32_t PRELOAD_NUM = 2;
105	static constexpr uint32_t N_BUFFER_M_BASIC_SIZE = 256;	107	static constexpr uint32_t N_BUFFER_M_BASIC_SIZE = 256;
	@@ -111,7 +113,6 @@ private:
111	static constexpr uint32_t SYNC_V1_C2_FLAG = 8;	113	static constexpr uint32_t SYNC_V1_C2_FLAG = 8;
112	static constexpr uint32_t SYNC_C2_V2_FLAG = 9;	114	static constexpr uint32_t SYNC_C2_V2_FLAG = 9;
113		115
114	- static constexpr uint64_t kvHeadNum = 1ULL;
115	static constexpr uint64_t headDim = 512ULL;	116	static constexpr uint64_t headDim = 512ULL;
116		117
117	static constexpr uint32_t dbWorkspaceRatio = PRELOAD_NUM;	118	static constexpr uint32_t dbWorkspaceRatio = PRELOAD_NUM;
	@@ -150,6 +151,7 @@ private:
150	GlobalTensor<int32_t> actualSeqLengthsQGm;	151	GlobalTensor<int32_t> actualSeqLengthsQGm;
151	GlobalTensor<int32_t> actualSeqLengthsKVGm;	152	GlobalTensor<int32_t> actualSeqLengthsKVGm;
152	GlobalTensor<int32_t> actualSeqLengthsCmpKVGm;	153	GlobalTensor<int32_t> actualSeqLengthsCmpKVGm;
		154	+ GlobalTensor<int32_t> cmpResidualKVGm;
153		155
154	// workspace	156	// workspace
155	GlobalTensor<MM1_OUT_T> mm1ResGm;	157	GlobalTensor<MM1_OUT_T> mm1ResGm;
	@@ -180,6 +182,8 @@ private:
180	RunInfo &info);	182	RunInfo &info);
181	__aicore__ inline int32_t GetActualSeqLenQ(uint32_t bIdx);	183	__aicore__ inline int32_t GetActualSeqLenQ(uint32_t bIdx);
182	__aicore__ inline int32_t GetActualSeqLenKV(uint32_t bIdx);	184	__aicore__ inline int32_t GetActualSeqLenKV(uint32_t bIdx);
		185	+ __aicore__ inline int32_t GetActualSeqLenCmpKV(uint32_t bIdx, int32_t actualOriS2Size);
		186	+ __aicore__ inline int32_t GetCmpMaskS2Size(uint32_t bIdx, int32_t actualOriS2Size, int32_t actualCmpS2Size);
183	__aicore__ inline void GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx);	187	__aicore__ inline void GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx);
184	// ================================Mm1==============================================	188	// ================================Mm1==============================================
185	__aicore__ inline void ComputeMm1(const RunInfo &info);	189	__aicore__ inline void ComputeMm1(const RunInfo &info);
	@@ -198,7 +202,9 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::InitTilingData()
198		202
199	// baseParams	203	// baseParams
200	constInfo.batchSize = tilingData->baseParams.batchSize;	204	constInfo.batchSize = tilingData->baseParams.batchSize;
201	- constInfo.qHeadNum = constInfo.gSize = tilingData->baseParams.nNumOfQInOneGroup;	205	+ constInfo.gSize = tilingData->baseParams.nNumOfQInOneGroup;
		206	+ constInfo.kvHeadNum = (tilingData->baseParams.kvHeadNum == 0) ? 1 : tilingData->baseParams.kvHeadNum;
		207	+ constInfo.qHeadNum = constInfo.gSize * constInfo.kvHeadNum;
202	constInfo.kvSeqSize = tilingData->baseParams.kvSeqSize;	208	constInfo.kvSeqSize = tilingData->baseParams.kvSeqSize;
203	constInfo.qSeqSize = tilingData->baseParams.qSeqSize;	209	constInfo.qSeqSize = tilingData->baseParams.qSeqSize;
204	constInfo.oriMaxBlockNumPerBatch = tilingData->baseParams.oriMaxBlockNumPerBatch;	210	constInfo.oriMaxBlockNumPerBatch = tilingData->baseParams.oriMaxBlockNumPerBatch;
	@@ -207,7 +213,6 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::InitTilingData()
207	constInfo.paOriBlockSize = tilingData->baseParams.oriBlockSize;	213	constInfo.paOriBlockSize = tilingData->baseParams.oriBlockSize;
208	constInfo.paCmpBlockSize = tilingData->baseParams.cmpBlockSize;	214	constInfo.paCmpBlockSize = tilingData->baseParams.cmpBlockSize;
209	constInfo.outputLayout = static_cast<SMLA_LAYOUT>(tilingData->baseParams.outputLayout);	215	constInfo.outputLayout = static_cast<SMLA_LAYOUT>(tilingData->baseParams.outputLayout);
210	- constInfo.kvHeadNum = kvHeadNum;
211	constInfo.headDim = headDim;	216	constInfo.headDim = headDim;
212	constInfo.oriMaskMode = tilingData->baseParams.oriMaskMode;	217	constInfo.oriMaskMode = tilingData->baseParams.oriMaskMode;
213	constInfo.oriKvStride0 = tilingData->baseParams.oriKvStride0;	218	constInfo.oriKvStride0 = tilingData->baseParams.oriKvStride0;
	@@ -216,6 +221,8 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::InitTilingData()
216		221
217	constInfo.actualLenDimsQ = tilingData->baseParams.actualLenDimsQ;	222	constInfo.actualLenDimsQ = tilingData->baseParams.actualLenDimsQ;
218	constInfo.actualLenDimsKV = tilingData->baseParams.actualLenDimsKV;	223	constInfo.actualLenDimsKV = tilingData->baseParams.actualLenDimsKV;
		224	+ constInfo.actualLenDimsCmpKV = tilingData->baseParams.actualLenDimsCmpKV;
		225	+ constInfo.cmpResidualKVSize = tilingData->baseParams.cmpResidualKVSize;
219	constInfo.returnSoftmaxLse = tilingData->baseParams.returnSoftmaxLse;	226	constInfo.returnSoftmaxLse = tilingData->baseParams.returnSoftmaxLse;
220	// innerSplitParams	227	// innerSplitParams
221	constInfo.mBaseSize = constInfo.gSize;	228	constInfo.mBaseSize = constInfo.gSize;
	@@ -234,6 +241,7 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::InitTilingData()
234	constInfo.sparseBlockSize = 1; // sparseBlockSize 固定为1	241	constInfo.sparseBlockSize = 1; // sparseBlockSize 固定为1
235	constInfo.cmpMaskMode = tilingData->cmpParams.cmpMaskMode;	242	constInfo.cmpMaskMode = tilingData->cmpParams.cmpMaskMode;
236	constInfo.cmpKvStride0 = tilingData->cmpParams.cmpKvStride0;	243	constInfo.cmpKvStride0 = tilingData->cmpParams.cmpKvStride0;
		244	+ constInfo.cmpSeqSize = tilingData->cmpParams.cmpKvSeqSize;
237	}	245	}
238		246
239	template <typename SMLAT>	247	template <typename SMLAT>
	@@ -265,7 +273,10 @@ SparseFlashMlaScfa<SMLAT>::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, _
265	{	273	{
266	if (constInfo.actualLenDimsKV != 0) {	274	if (constInfo.actualLenDimsKV != 0) {
267	actualSeqLengthsKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsKV, constInfo.actualLenDimsKV);	275	actualSeqLengthsKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsKV, constInfo.actualLenDimsKV);
268	- actualSeqLengthsCmpKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsCmpKV, constInfo.actualLenDimsKV);	276	+ }
		277	+ if (constInfo.actualLenDimsCmpKV != 0) {
		278	+ actualSeqLengthsCmpKVGm.SetGlobalBuffer(
		279	+ (__gm__ int32_t *)actualSeqLengthsCmpKV, constInfo.actualLenDimsCmpKV);
269	}	280	}
270	if (constInfo.actualLenDimsQ != 0) {	281	if (constInfo.actualLenDimsQ != 0) {
271	actualSeqLengthsQGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsQ, constInfo.actualLenDimsQ);	282	actualSeqLengthsQGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsQ, constInfo.actualLenDimsQ);
	@@ -282,22 +293,24 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::InitAllZeroOutput(uint32_t bId
282	uint32_t tBase = actualSeqLengthsQGm.GetValue(bIdx);	293	uint32_t tBase = actualSeqLengthsQGm.GetValue(bIdx);
283	uint32_t s1Count = tempLoopInfo.actS1Size;	294	uint32_t s1Count = tempLoopInfo.actS1Size;
284		295
285	- uint64_t attenOutOffset = (tBase + s1Idx) * kvHeadNum * constInfo.gSize * headDim + // T轴、s1轴偏移	296	+ uint64_t attenOutOffset = (tBase + s1Idx) * constInfo.kvHeadNum * constInfo.gSize *
286	- n2Idx * constInfo.gSize * headDim; // N2轴偏移	297	+ constInfo.headDim +
		298	+ n2Idx * constInfo.gSize * constInfo.headDim;
287	uint64_t lseOffset = (tBase + s1Idx) * constInfo.gSize + // T轴、s1轴偏移	299	uint64_t lseOffset = (tBase + s1Idx) * constInfo.gSize + // T轴、s1轴偏移
288	n2Idx * constInfo.qSeqSize * constInfo.gSize; // N2轴偏移	300	n2Idx * constInfo.qSeqSize * constInfo.gSize; // N2轴偏移
289	- matmul::InitOutput<OUT_T>(attentionOutGm[attenOutOffset], constInfo.gSize * headDim, 0);	301	+ matmul::InitOutput<OUT_T>(attentionOutGm[attenOutOffset], constInfo.gSize * constInfo.headDim, 0);
290	if (constInfo.returnSoftmaxLse) {	302	if (constInfo.returnSoftmaxLse) {
291	matmul::InitOutput<T>(softmaxLseGm[lseOffset], constInfo.gSize, 0);	303	matmul::InitOutput<T>(softmaxLseGm[lseOffset], constInfo.gSize, 0);
292	}	304	}
293	} else if (constInfo.outputLayout == SMLA_LAYOUT::BSND) {	305	} else if (constInfo.outputLayout == SMLA_LAYOUT::BSND) {
294	- uint64_t attenOutOffset = bIdx * constInfo.qSeqSize * kvHeadNum * constInfo.gSize * headDim +	306	+ uint64_t attenOutOffset = bIdx * constInfo.qSeqSize * constInfo.kvHeadNum * constInfo.gSize *
295	- s1Idx * kvHeadNum * constInfo.gSize * headDim + // B轴、S1轴偏移	307	+ constInfo.headDim +
296	- n2Idx * constInfo.gSize * headDim; // N2轴偏移	308	+ s1Idx * constInfo.kvHeadNum * constInfo.gSize * constInfo.headDim +
		309	+ n2Idx * constInfo.gSize * constInfo.headDim;
297	uint64_t lseOffset = bIdx * constInfo.qSeqSize * constInfo.kvHeadNum * constInfo.gSize + // B轴偏移	310	uint64_t lseOffset = bIdx * constInfo.qSeqSize * constInfo.kvHeadNum * constInfo.gSize + // B轴偏移
298	n2Idx * constInfo.qSeqSize * constInfo.gSize + // N2轴偏移	311	n2Idx * constInfo.qSeqSize * constInfo.gSize + // N2轴偏移
299	s1Idx * constInfo.gSize; // S1轴偏移	312	s1Idx * constInfo.gSize; // S1轴偏移
300	- matmul::InitOutput<OUT_T>(attentionOutGm[attenOutOffset], constInfo.gSize * headDim, 0);	313	+ matmul::InitOutput<OUT_T>(attentionOutGm[attenOutOffset], constInfo.gSize * constInfo.headDim, 0);
301	if (constInfo.returnSoftmaxLse) {	314	if (constInfo.returnSoftmaxLse) {
302	matmul::InitOutput<T>(softmaxLseGm[lseOffset], constInfo.gSize, 0);	315	matmul::InitOutput<T>(softmaxLseGm[lseOffset], constInfo.gSize, 0);
303	}	316	}
	@@ -348,16 +361,56 @@ __aicore__ inline int32_t SparseFlashMlaScfa<SMLAT>::GetActualSeqLenKV(uint32_t
348	}	361	}
349	return actualSeqLengthsKVGm.GetValue(bIdx);	362	return actualSeqLengthsKVGm.GetValue(bIdx);
350	} else if constexpr(KV_LAYOUT_T == SMLA_LAYOUT::BSND) {	363	} else if constexpr(KV_LAYOUT_T == SMLA_LAYOUT::BSND) {
		364	+ tempLoopInfo.actualSeqKVPrefixSum = static_cast<uint64_t>(bIdx * constInfo.kvSeqSize);
		365	+ if (constInfo.actualLenDimsKV != 0) {
		366	+ return actualSeqLengthsKVGm.GetValue(bIdx);
		367	+ }
351	return static_cast<int32_t>(constInfo.kvSeqSize);	368	return static_cast<int32_t>(constInfo.kvSeqSize);
352	} else if constexpr(KV_LAYOUT_T == SMLA_LAYOUT::TND) {	369	} else if constexpr(KV_LAYOUT_T == SMLA_LAYOUT::TND) {
353	int32_t actualSeqKVPrefixSum = actualSeqLengthsKVGm.GetValue(bIdx);	370	int32_t actualSeqKVPrefixSum = actualSeqLengthsKVGm.GetValue(bIdx);
354	int32_t actualSeqKVNextSum = actualSeqLengthsKVGm.GetValue(bIdx + 1);	371	int32_t actualSeqKVNextSum = actualSeqLengthsKVGm.GetValue(bIdx + 1);
355	- tempLoopInfo.actualSeqCmpKVPrefixSum = actualSeqLengthsCmpKVGm.GetValue(bIdx);
356	tempLoopInfo.actualSeqKVPrefixSum = actualSeqKVPrefixSum;	372	tempLoopInfo.actualSeqKVPrefixSum = actualSeqKVPrefixSum;
357	return actualSeqKVNextSum - actualSeqKVPrefixSum;	373	return actualSeqKVNextSum - actualSeqKVPrefixSum;
358	}	374	}
359	}	375	}
360		376
		377	+template <typename SMLAT>
		378	+__aicore__ inline int32_t SparseFlashMlaScfa<SMLAT>::GetActualSeqLenCmpKV(uint32_t bIdx, int32_t actualOriS2Size)
		379	+{
		380	+ (void)actualOriS2Size;
		381	+ if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::TND) {
		382	+ int32_t actualSeqCmpKVPrefixSum = actualSeqLengthsCmpKVGm.GetValue(bIdx);
		383	+ int32_t actualSeqCmpKVNextSum = actualSeqLengthsCmpKVGm.GetValue(bIdx + 1);
		384	+ tempLoopInfo.actualSeqCmpKVPrefixSum = actualSeqCmpKVPrefixSum;
		385	+ return actualSeqCmpKVNextSum - actualSeqCmpKVPrefixSum;
		386	+ } else if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND) {
		387	+ tempLoopInfo.actualSeqCmpKVPrefixSum = static_cast<uint64_t>(bIdx * constInfo.cmpSeqSize);
		388	+ if (constInfo.actualLenDimsCmpKV != 0) {
		389	+ return actualSeqLengthsCmpKVGm.GetValue(bIdx);
		390	+ }
		391	+ return static_cast<int32_t>(constInfo.cmpSeqSize);
		392	+ } else {
		393	+ tempLoopInfo.actualSeqCmpKVPrefixSum = static_cast<uint64_t>(bIdx * constInfo.cmpSeqSize);
		394	+ if (constInfo.actualLenDimsCmpKV != 0) {
		395	+ return actualSeqLengthsCmpKVGm.GetValue(bIdx);
		396	+ }
		397	+ return (constInfo.cmpSeqSize != 0) ? static_cast<int32_t>(constInfo.cmpSeqSize) :
		398	+ actualOriS2Size / static_cast<int32_t>(constInfo.cmpRatio);
		399	+ }
		400	+}
		401	+
		402	+template <typename SMLAT>
		403	+__aicore__ inline int32_t SparseFlashMlaScfa<SMLAT>::GetCmpMaskS2Size(uint32_t bIdx, int32_t actualOriS2Size,
		404	+ int32_t actualCmpS2Size)
		405	+{
		406	+ (void)actualOriS2Size;
		407	+ int32_t residual = 0;
		408	+ if (constInfo.cmpResidualKVSize != 0) {
		409	+ residual = cmpResidualKVGm.GetValue(bIdx);
		410	+ }
		411	+ return actualCmpS2Size * static_cast<int32_t>(constInfo.cmpRatio) + residual;
		412	+}
		413	+
361	template <typename SMLAT>	414	template <typename SMLAT>
362	__aicore__ inline void SparseFlashMlaScfa<SMLAT>::GetSparseActualSeqLen()	415	__aicore__ inline void SparseFlashMlaScfa<SMLAT>::GetSparseActualSeqLen()
363	{	416	{
	@@ -370,7 +423,8 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::GetSparseActualSeqLen()
370		423
371	// 对于cmp部分还有top k, tempLoopInfo.actS2Size只针对cmp	424	// 对于cmp部分还有top k, tempLoopInfo.actS2Size只针对cmp
372	int32_t thresHold = (tempLoopInfo.cmpMaskRight + tempLoopInfo.s1EndIdx + 1) / constInfo.cmpRatio;	425	int32_t thresHold = (tempLoopInfo.cmpMaskRight + tempLoopInfo.s1EndIdx + 1) / constInfo.cmpRatio;
373	- tempLoopInfo.actCmpS2Size = Min(constInfo.sparseBlockCount * constInfo.sparseBlockSize, thresHold);	426	+ tempLoopInfo.actCmpS2Size = Min(tempLoopInfo.actCmpS2Size,
		427	+ Min(constInfo.sparseBlockCount * constInfo.sparseBlockSize, Max(thresHold, 0)));
374	}	428	}
375		429
376	template <typename SMLAT>	430	template <typename SMLAT>
	@@ -391,7 +445,8 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::Init(
391	__gm__ uint8_t query, __gm__ uint8_t oriKV, __gm__ uint8_t cmpKV, __gm__ uint8_t cmpSparseIndices,	445	__gm__ uint8_t query, __gm__ uint8_t oriKV, __gm__ uint8_t cmpKV, __gm__ uint8_t cmpSparseIndices,
392	__gm__ uint8_t oriBlockTable, __gm__ uint8_t cmpBlockTable, __gm__ uint8_t *cuSeqlensQ,	446	__gm__ uint8_t oriBlockTable, __gm__ uint8_t cmpBlockTable, __gm__ uint8_t *cuSeqlensQ,
393	__gm__ uint8_t* cuSeqlensKV, __gm__ uint8_t cuSeqlensCmpKV, __gm__ uint8_t seqUsedQ,	447	__gm__ uint8_t* cuSeqlensKV, __gm__ uint8_t cuSeqlensCmpKV, __gm__ uint8_t seqUsedQ,
394	- __gm__ uint8_t seqUsedKV, __gm__ uint8_t sinks, __gm__ uint8_t metadata, __gm__ uint8_t attentionOut, __gm__ uint8_t *softmaxLse,	448	+ __gm__ uint8_t seqUsedKV, __gm__ uint8_t seqUsedCmpKV, __gm__ uint8_t *cmpResidualKV,
		449	+ __gm__ uint8_t sinks, __gm__ uint8_t metadata, __gm__ uint8_t attentionOut, __gm__ uint8_t softmaxLse,
395	__gm__ uint8_t workspace, const SparseFlashMlaTilingData __restrict tiling, __gm__ uint8_t *gmTiling,	450	__gm__ uint8_t workspace, const SparseFlashMlaTilingData __restrict tiling, __gm__ uint8_t *gmTiling,
396	TPipe *tPipe)	451	TPipe *tPipe)
397	{	452	{
	@@ -412,9 +467,9 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::Init(
412	InitActualSeqLen(seqUsedQ, cuSeqlensKV, cuSeqlensCmpKV);	467	InitActualSeqLen(seqUsedQ, cuSeqlensKV, cuSeqlensCmpKV);
413	} else if ((KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND \|\| KV_LAYOUT_T == SMLA_LAYOUT::BSND)	468	} else if ((KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND \|\| KV_LAYOUT_T == SMLA_LAYOUT::BSND)
414	&& LAYOUT_T == SMLA_LAYOUT::TND) {	469	&& LAYOUT_T == SMLA_LAYOUT::TND) {
415	- InitActualSeqLen(cuSeqlensQ, seqUsedKV);	470	+ InitActualSeqLen(cuSeqlensQ, seqUsedKV, seqUsedCmpKV);
416	} else if ((KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND \|\| KV_LAYOUT_T == SMLA_LAYOUT::BSND)) {	471	} else if ((KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND \|\| KV_LAYOUT_T == SMLA_LAYOUT::BSND)) {
417	- InitActualSeqLen(seqUsedQ, seqUsedKV);	472	+ InitActualSeqLen(seqUsedQ, seqUsedKV, seqUsedCmpKV);
418	}	473	}
419		474
420	metadataGm.SetGlobalBuffer((__gm__ uint32_t *)metadata);	475	metadataGm.SetGlobalBuffer((__gm__ uint32_t *)metadata);
	@@ -425,6 +480,9 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::Init(
425	queryGm.SetGlobalBuffer((__gm__ Q_T *)query);	480	queryGm.SetGlobalBuffer((__gm__ Q_T *)query);
426	oriKvGm.SetGlobalBuffer((__gm__ KV_T *)oriKV);	481	oriKvGm.SetGlobalBuffer((__gm__ KV_T *)oriKV);
427	cmpKvGm.SetGlobalBuffer((__gm__ KV_T *)cmpKV);	482	cmpKvGm.SetGlobalBuffer((__gm__ KV_T *)cmpKV);
		483	+ if (constInfo.cmpResidualKVSize != 0) {
		484	+ cmpResidualKVGm.SetGlobalBuffer((__gm__ int32_t *)cmpResidualKV, constInfo.cmpResidualKVSize);
		485	+ }
428		486
429	if (sinks != nullptr) {	487	if (sinks != nullptr) {
430	sinksGm.SetGlobalBuffer((__gm__ SINKS_T *)sinks);	488	sinksGm.SetGlobalBuffer((__gm__ SINKS_T *)sinks);
	@@ -552,7 +610,9 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::CalcParams(uint32_t loop, uint
552	uint64_t tndBIdxOffsetForCmpKV = tempLoopInfo.actualSeqCmpKVPrefixSum * constInfo.kvHeadNum * constInfo.headDim;	610	uint64_t tndBIdxOffsetForCmpKV = tempLoopInfo.actualSeqCmpKVPrefixSum * constInfo.kvHeadNum * constInfo.headDim;
553		611
554	if (info.isFirstSInnerLoop) {	612	if (info.isFirstSInnerLoop) {
555	- tensorACoreOffset = tndBIdxOffsetForQ + info.gS1Idx * constInfo.headDim;	613	+ uint64_t s1HeadOffset = (info.gS1Idx / constInfo.gSize) * constInfo.qHeadNum;
		614	+ uint64_t qHeadOffset = info.n2Idx * constInfo.gSize + info.gS1Idx % constInfo.gSize;
		615	+ tensorACoreOffset = tndBIdxOffsetForQ + (s1HeadOffset + qHeadOffset) * constInfo.headDim;
556	tensorBCoreOffset = tndBIdxOffsetForKV + info.n2Idx * constInfo.headDim; // 当前为PA场景，该变量失效	616	tensorBCoreOffset = tndBIdxOffsetForKV + info.n2Idx * constInfo.headDim; // 当前为PA场景，该变量失效
557	tensorCmpBCoreOffset = tndBIdxOffsetForCmpKV + info.n2Idx * constInfo.headDim;	617	tensorCmpBCoreOffset = tndBIdxOffsetForCmpKV + info.n2Idx * constInfo.headDim;
558	if constexpr (LAYOUT_T == SMLA_LAYOUT::BSND) { // B,S1,N2 K	618	if constexpr (LAYOUT_T == SMLA_LAYOUT::BSND) { // B,S1,N2 K
	@@ -594,18 +654,22 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::CalcParams(uint32_t loop, uint
594	}	654	}
595	info.s2StartPoint = 0;	655	info.s2StartPoint = 0;
596	info.cmpS2IdLimit = (tempLoopInfo.cmpMaskRight + tempLoopInfo.s1EndIdx + 1) / constInfo.cmpRatio;	656	info.cmpS2IdLimit = (tempLoopInfo.cmpMaskRight + tempLoopInfo.s1EndIdx + 1) / constInfo.cmpRatio;
		657	+ if constexpr (HEAD_RATIO_ONE) {
		658	+ info.v0S2Start = static_cast<int32_t>(s2Offset);
		659	+ info.v0S2DealSize = static_cast<int32_t>(info.actualSingleProcessSInnerSize);
		660	+ } else {
		661	+ info.v0S2Start = 0;
		662	+ if (s2LoopIdx + 1 == tempLoopInfo.s2LoopTimes && s2LoopIdx == 2) {
		663	+ info.v0S2Start = 512;
		664	+ }
		665	+ info.v0S2DealSize = 512;
		666	+ }
597	}	667	}
598		668
599	info.actualSingleProcessSInnerSizeAlign = SMLAAlign(info.actualSingleProcessSInnerSize, SMLAVectorBlock<SMLAT>::BYTE_BLOCK);	669	info.actualSingleProcessSInnerSizeAlign = SMLAAlign(info.actualSingleProcessSInnerSize, SMLAVectorBlock<SMLAT>::BYTE_BLOCK);
600	if (info.isOriOnly) {	670	if (info.isOriOnly) {
601	info.v0S2Start = 0;	671	info.v0S2Start = 0;
602	info.v0S2DealSize = 0;	672	info.v0S2DealSize = 0;
603	- } else {
604	- info.v0S2Start = 0;
605	- if (s2LoopIdx + 1 == tempLoopInfo.s2LoopTimes && s2LoopIdx == 2) { // tail
606	- info.v0S2Start = 512;
607	- }
608	- info.v0S2DealSize = 512;
609	}	673	}
610	}	674	}
611		675
	@@ -619,6 +683,11 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::ComputeMm1(const RunInfo &info
619	mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize;	683	mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize;
620	mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail;	684	mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail;
621	cubeBlock.ComputeMm1(info, mSplitInfo);	685	cubeBlock.ComputeMm1(info, mSplitInfo);
		686	+ if constexpr (HEAD_RATIO_ONE) {
		687	+ event_t eventIdFixWait = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::FIX_M));
		688	+ SetFlag<HardEvent::FIX_M>(eventIdFixWait);
		689	+ WaitFlag<HardEvent::FIX_M>(eventIdFixWait);
		690	+ }
622	CrossCoreSetFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_FIX>(constInfo.syncC1V1);	691	CrossCoreSetFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_FIX>(constInfo.syncC1V1);
623	}	692	}
624	}	693	}
	@@ -632,8 +701,17 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::ComputeMm2(const RunInfo &info
632	MSplitInfo mSplitInfo;	701	MSplitInfo mSplitInfo;
633	mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize;	702	mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize;
634	mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail;	703	mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail;
635	- CrossCoreWaitFlag(constInfo.syncV1C2);	704	+ if constexpr (HEAD_RATIO_ONE) {
		705	+ CrossCoreWaitFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_MTE3>(constInfo.syncV1C2);
		706	+ } else {
		707	+ CrossCoreWaitFlag(constInfo.syncV1C2);
		708	+ }
636	cubeBlock.ComputeMm2(info, mSplitInfo);	709	cubeBlock.ComputeMm2(info, mSplitInfo);
		710	+ if constexpr (HEAD_RATIO_ONE) {
		711	+ event_t eventIdFixWait = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::FIX_M));
		712	+ SetFlag<HardEvent::FIX_M>(eventIdFixWait);
		713	+ WaitFlag<HardEvent::FIX_M>(eventIdFixWait);
		714	+ }
637	CrossCoreSetFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_FIX>(constInfo.syncC2V2);	715	CrossCoreSetFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_FIX>(constInfo.syncC2V2);
638	}	716	}
639	}	717	}
	@@ -662,8 +740,8 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::Process()
662	template <typename SMLAT>	740	template <typename SMLAT>
663	__aicore__ inline void SparseFlashMlaScfa<SMLAT>::GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx)	741	__aicore__ inline void SparseFlashMlaScfa<SMLAT>::GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx)
664	{	742	{
665	- bIdx = bN2Idx / kvHeadNum;	743	+ bIdx = bN2Idx / constInfo.kvHeadNum;
666	- n2Idx = bN2Idx % kvHeadNum;	744	+ n2Idx = bN2Idx % constInfo.kvHeadNum;
667	}	745	}
668		746
669	template <typename SMLAT>	747	template <typename SMLAT>
	@@ -704,6 +782,7 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::ProcessBalance()
704	gS1LoopEnd = isS1ZeroAndLastBatch ? gS1LoopEnd + 1 : gS1LoopEnd;	782	gS1LoopEnd = isS1ZeroAndLastBatch ? gS1LoopEnd + 1 : gS1LoopEnd;
705	for (uint32_t gS1LoopIdx = constInfo.gS1Start; gS1LoopIdx < gS1LoopEnd; gS1LoopIdx++) {	783	for (uint32_t gS1LoopIdx = constInfo.gS1Start; gS1LoopIdx < gS1LoopEnd; gS1LoopIdx++) {
706	tempLoopInfo.actOriS2Size = GetActualSeqLenKV(tempLoopInfo.bIdx);	784	tempLoopInfo.actOriS2Size = GetActualSeqLenKV(tempLoopInfo.bIdx);
		785	+ tempLoopInfo.actCmpS2Size = GetActualSeqLenCmpKV(tempLoopInfo.bIdx, tempLoopInfo.actOriS2Size);
707	// 计算需要的数据, 避免重复计算	786	// 计算需要的数据, 避免重复计算
708	tempLoopInfo.gS1Idx = gS1LoopIdx * constInfo.mBaseSize;	787	tempLoopInfo.gS1Idx = gS1LoopIdx * constInfo.mBaseSize;
709	tempLoopInfo.s1StartIdx = tempLoopInfo.gS1Idx / constInfo.gSize;	788	tempLoopInfo.s1StartIdx = tempLoopInfo.gS1Idx / constInfo.gSize;
	@@ -716,7 +795,9 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::ProcessBalance()
716	tempLoopInfo.oriMaskLeft = Max(tempLoopInfo.actOriS2Size - tempLoopInfo.actS1Size +	795	tempLoopInfo.oriMaskLeft = Max(tempLoopInfo.actOriS2Size - tempLoopInfo.actS1Size +
717	static_cast<int32_t>(tempLoopInfo.s1EndIdx) - constInfo.oriWinLeft,	796	static_cast<int32_t>(tempLoopInfo.s1EndIdx) - constInfo.oriWinLeft,
718	0);	797	0);
719	- tempLoopInfo.cmpMaskRight = tempLoopInfo.actOriS2Size - tempLoopInfo.actS1Size;	798	+ int32_t cmpMaskS2Size = GetCmpMaskS2Size(
		799	+ tempLoopInfo.bIdx, tempLoopInfo.actOriS2Size, tempLoopInfo.actCmpS2Size);
		800	+ tempLoopInfo.cmpMaskRight = cmpMaskS2Size - tempLoopInfo.actS1Size;
720	GetSparseActualSeqLen();	801	GetSparseActualSeqLen();
721	UpdateInnerLoopCond();	802	UpdateInnerLoopCond();
722		803
	@@ -790,38 +871,77 @@ __aicore__ inline void SparseFlashMlaScfa<SMLAT>::PreloadPipeline(uint32_t loop,
790	RunInfo &extraInfo1 = extraInfo[(loop + 1) % SMLA_PRELOAD_TASK_CACHE_SIZE]; // 上两轮任务	871	RunInfo &extraInfo1 = extraInfo[(loop + 1) % SMLA_PRELOAD_TASK_CACHE_SIZE]; // 上两轮任务
791		872
792	CalcParams(loop, cmpLoop, s2Start, s2LoopIdx, extraInfo0);	873	CalcParams(loop, cmpLoop, s2Start, s2LoopIdx, extraInfo0);
793	- if (extraInfo0.isValid) {	874	+ if constexpr (!HEAD_RATIO_ONE) {
794	- if ASCEND_IS_AIC {	875	+ if (extraInfo0.isValid) {
795	- if (!extraInfo0.isOriOnly) {	876	+ if ASCEND_IS_AIC {
796	- CrossCoreWaitFlag(constInfo.syncV0C1);	877	+ if (!extraInfo0.isOriOnly) {
797	- }	878	+ CrossCoreWaitFlag(constInfo.syncV0C1);
798	- ComputeMm1(extraInfo0);	879	+ }
799	- } else {	880	+ ComputeMm1(extraInfo0);
800	- if (extraInfo0.isFirstSInnerLoop) {	881	+ } else {
801	- CrossCoreWaitFlag(3);	882	+ if (extraInfo0.isFirstSInnerLoop) {
802	- }	883	+ CrossCoreWaitFlag(3);
803	- vectorBlock.ProcessVec0L(extraInfo0);	884	+ }
804	- if (!extraInfo0.isOriOnly) {	885	+ vectorBlock.ProcessVec0L(extraInfo0);
805	- CrossCoreSetFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_MTE3>(constInfo.syncV0C1);	886	+ if (!extraInfo0.isOriOnly) {
		887	+ CrossCoreSetFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_MTE3>(constInfo.syncV0C1);
		888	+ }
806	}	889	}
807	}	890	}
808	- }	891	+ if (extraInfo2.isValid) {
809	- if (extraInfo2.isValid) {	892	+ if ASCEND_IS_AIV {
810	- if ASCEND_IS_AIV {	893	+ vectorBlock.ProcessVec1L(extraInfo2);
811	- vectorBlock.ProcessVec1L(extraInfo2);	894	+ }
812	- }	895	+ if ASCEND_IS_AIC {
813	- if ASCEND_IS_AIC {	896	+ ComputeMm2(extraInfo2);
814	- ComputeMm2(extraInfo2);	897	+ if (extraInfo2.isLastS2Loop) {
815	- if (extraInfo2.isLastS2Loop) {	898	+ CrossCoreSetFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_MTE2>(3);
816	- CrossCoreSetFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_MTE2>(3);	899	+ }
817	}	900	}
818	}	901	}
819	- }	902	+ if (extraInfo1.isValid) {
820	- if (extraInfo1.isValid) {	903	+ if ASCEND_IS_AIV {
821	- if ASCEND_IS_AIV {	904	+ vectorBlock.ProcessVec2L(extraInfo1);
822	- vectorBlock.ProcessVec2L(extraInfo1);	905	+ }
		906	+ extraInfo1.isValid = false;
		907	+ }
		908	+ } else {
		909	+ if (extraInfo0.isValid) {
		910	+ if ASCEND_IS_AIV {
		911	+ if (extraInfo0.isFirstSInnerLoop) {
		912	+ CrossCoreWaitFlag(3);
		913	+ }
		914	+ vectorBlock.ProcessVec0L(extraInfo0);
		915	+ if (!extraInfo0.isOriOnly) {
		916	+ CrossCoreSetFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_MTE3>(constInfo.syncV0C1);
		917	+ }
		918	+ }
		919	+ }
		920	+ if (extraInfo2.isValid) {
		921	+ if ASCEND_IS_AIV {
		922	+ vectorBlock.ProcessVec1L(extraInfo2);
		923	+ }
		924	+ if ASCEND_IS_AIC {
		925	+ ComputeMm2(extraInfo2);
		926	+ if (extraInfo2.isLastS2Loop) {
		927	+ CrossCoreSetFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_MTE2>(3);
		928	+ }
		929	+ }
		930	+ }
		931	+ if (extraInfo1.isValid) {
		932	+ if ASCEND_IS_AIV {
		933	+ vectorBlock.ProcessVec2L(extraInfo1);
		934	+ }
		935	+ extraInfo1.isValid = false;
		936	+ }
		937	+ if (extraInfo0.isValid) {
		938	+ if ASCEND_IS_AIC {
		939	+ if (!extraInfo0.isOriOnly) {
		940	+ CrossCoreWaitFlag<ConstInfo::SMLA_SYNC_MODE2, PIPE_MTE3>(constInfo.syncV0C1);
		941	+ }
		942	+ ComputeMm1(extraInfo0);
		943	+ }
823	}	944	}
824	- extraInfo1.isValid = false;
825	}	945	}
826	}	946	}
827		947

Mattention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_swa_block_cube.h+20-14

	@@ -413,7 +413,8 @@ __aicore__ inline void SWACubeBlock<SMLAT>::ComputeMm1(const RunInfo &info, cons
413		413
414	uint32_t headStride = constInfo.headDim;	414	uint32_t headStride = constInfo.headDim;
415	uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim;	415	uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim;
416	- uint32_t batchStride = constInfo.kvSeqSize * seqStride;	416	+ uint64_t batchStride = (constInfo.oriKvStride0 == 0) ?
		417	+ static_cast<uint64_t>(constInfo.kvSeqSize) * seqStride : constInfo.oriKvStride0;
417		418
418	uint64_t curS2 = (uint64_t)info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint;	419	uint64_t curS2 = (uint64_t)info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint;
419	uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride + \	420	uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride + \
	@@ -507,11 +508,7 @@ __aicore__ inline void SWACubeBlock<SMLAT>::ComputeMm1(const RunInfo &info, cons
507	shape.copyRowNum = copyRowCnt;	508	shape.copyRowNum = copyRowCnt;
508	shape.copyRowNumAlign = nL1SizeAlign;	509	shape.copyRowNumAlign = nL1SizeAlign;
509	kTensor = bL1Tensor[copyFinishRowCnt * 16 + oriSizeCur * 16];	510	kTensor = bL1Tensor[copyFinishRowCnt * 16 + oriSizeCur * 16];
510	- if (nL1 == 0) {	511	+ DataCopyPA<KV_T>(kTensor, cmpKvGm, cmpBlockTableGm, shape, startPos);
511	- DataCopyPA<KV_T>(kTensor, oriKvGm, oriBlockTableGm, shape, startPos);
512	- } else {
513	- DataCopyPA<KV_T>(kTensor, cmpKvGm, cmpBlockTableGm, shape, startPos);
514	- }
515	// 更新循环变量	512	// 更新循环变量
516	copyFinishRowCnt += copyRowCnt;	513	copyFinishRowCnt += copyRowCnt;
517	curS2Offset += copyRowCnt;	514	curS2Offset += copyRowCnt;
	@@ -531,7 +528,9 @@ __aicore__ inline void SWACubeBlock<SMLAT>::ComputeMm1(const RunInfo &info, cons
531	uint32_t headStride = constInfo.headDim;	528	uint32_t headStride = constInfo.headDim;
532	uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim;	529	uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim;
533	if (oriSizeCur > 0) {	530	if (oriSizeCur > 0) {
534	- uint32_t batchStride = constInfo.kvSeqSize * seqStride;	531	+ nd2nzPara.nValue = oriSizeCur;
		532	+ uint64_t batchStride = (constInfo.oriKvStride0 == 0) ?
		533	+ static_cast<uint64_t>(constInfo.kvSeqSize) * seqStride : constInfo.oriKvStride0;
535	uint64_t curS2 = (uint64_t)info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint + \	534	uint64_t curS2 = (uint64_t)info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint + \
536	nL1 * N_SPLIT_SIZE;	535	nL1 * N_SPLIT_SIZE;
537	uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride + \	536	uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride + \
	@@ -539,7 +538,9 @@ __aicore__ inline void SWACubeBlock<SMLAT>::ComputeMm1(const RunInfo &info, cons
539	DataCopy(bL1Tensor, oriKvGm[offset], nd2nzPara);	538	DataCopy(bL1Tensor, oriKvGm[offset], nd2nzPara);
540	}	539	}
541	if (cmpSizeCur > 0) {	540	if (cmpSizeCur > 0) {
542	- uint32_t batchStride = constInfo.kvSeqSize / constInfo.cmpRatio * seqStride;	541	+ nd2nzPara.nValue = cmpSizeCur;
		542	+ uint64_t batchStride = (constInfo.cmpKvStride0 == 0) ?
		543	+ static_cast<uint64_t>(constInfo.cmpSeqSize) * seqStride : constInfo.cmpKvStride0;
543	uint32_t cmpMixsizeCur = N_SPLIT_SIZE - info.actualSingleProcessSInnerOriSize % N_SPLIT_SIZE;	544	uint32_t cmpMixsizeCur = N_SPLIT_SIZE - info.actualSingleProcessSInnerOriSize % N_SPLIT_SIZE;
544	uint32_t oriOnlyLoopTimes = info.actualSingleProcessSInnerOriSize / N_SPLIT_SIZE;	545	uint32_t oriOnlyLoopTimes = info.actualSingleProcessSInnerOriSize / N_SPLIT_SIZE;
545	uint32_t cmpLoopTimes = nL1 - oriOnlyLoopTimes;	546	uint32_t cmpLoopTimes = nL1 - oriOnlyLoopTimes;
	@@ -625,7 +626,8 @@ __aicore__ inline void SWACubeBlock<SMLAT>::ComputeMm1(const RunInfo &info, cons
625		626
626	uint32_t headStride = constInfo.headDim;	627	uint32_t headStride = constInfo.headDim;
627	uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim;	628	uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim;
628	- uint32_t batchStride = constInfo.kvSeqSize / constInfo.cmpRatio * seqStride;	629	+ uint64_t batchStride = (constInfo.cmpKvStride0 == 0) ?
		630	+ static_cast<uint64_t>(constInfo.cmpSeqSize) * seqStride : constInfo.cmpKvStride0;
629		631
630	uint64_t curS2 = static_cast<uint64_t>(info.relativeS2Idx) * constInfo.s2BaseSize + \	632	uint64_t curS2 = static_cast<uint64_t>(info.relativeS2Idx) * constInfo.s2BaseSize + \
631	info.s2StartPoint + nL1 * N_SPLIT_SIZE;	633	info.s2StartPoint + nL1 * N_SPLIT_SIZE;
	@@ -844,7 +846,8 @@ __aicore__ inline void SWACubeBlock<SMLAT>::ComputeMm2(const RunInfo &info, cons
844		846
845	uint32_t headStride = constInfo.headDim;	847	uint32_t headStride = constInfo.headDim;
846	uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim;	848	uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim;
847	- uint32_t batchStride = constInfo.kvSeqSize * seqStride;	849	+ uint64_t batchStride = (constInfo.oriKvStride0 == 0) ?
		850	+ static_cast<uint64_t>(constInfo.kvSeqSize) * seqStride : constInfo.oriKvStride0;
848		851
849	uint64_t curS2 = (uint64_t)info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint + \	852	uint64_t curS2 = (uint64_t)info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint + \
850	kL1 * K_L0_SPLIT_SIZE;	853	kL1 * K_L0_SPLIT_SIZE;
	@@ -962,7 +965,8 @@ __aicore__ inline void SWACubeBlock<SMLAT>::ComputeMm2(const RunInfo &info, cons
962	if (oriSizeCur > 0) {	965	if (oriSizeCur > 0) {
963	subvTensor = bL1Tensor[(kL1 - kOffset) * K_L0_SPLIT_SIZE * N_SPLIT_SIZE];	966	subvTensor = bL1Tensor[(kL1 - kOffset) * K_L0_SPLIT_SIZE * N_SPLIT_SIZE];
964	nd2nzPara.nValue = oriSizeCur;	967	nd2nzPara.nValue = oriSizeCur;
965	- uint32_t batchStride = constInfo.kvSeqSize * seqStride;	968	+ uint64_t batchStride = (constInfo.oriKvStride0 == 0) ?
		969	+ static_cast<uint64_t>(constInfo.kvSeqSize) * seqStride : constInfo.oriKvStride0;
966	uint64_t curS2 = (uint64_t)info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint +	970	uint64_t curS2 = (uint64_t)info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint +
967	kL1 * K_L0_SPLIT_SIZE;	971	kL1 * K_L0_SPLIT_SIZE;
968	uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride +	972	uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride +
	@@ -972,7 +976,8 @@ __aicore__ inline void SWACubeBlock<SMLAT>::ComputeMm2(const RunInfo &info, cons
972	if (cmpSizeCur > 0) {	976	if (cmpSizeCur > 0) {
973	nd2nzPara.nValue = cmpSizeCur;	977	nd2nzPara.nValue = cmpSizeCur;
974	subvTensor = bL1Tensor[(kL1 - kOffset) * K_L0_SPLIT_SIZE * N_SPLIT_SIZE + oriSizeCur * 16];	978	subvTensor = bL1Tensor[(kL1 - kOffset) * K_L0_SPLIT_SIZE * N_SPLIT_SIZE + oriSizeCur * 16];
975	- uint32_t batchStride = constInfo.kvSeqSize / constInfo.cmpRatio * seqStride;	979	+ uint64_t batchStride = (constInfo.cmpKvStride0 == 0) ?
		980	+ static_cast<uint64_t>(constInfo.cmpSeqSize) * seqStride : constInfo.cmpKvStride0;
976	uint32_t cmpMixsizeCur = K_L0_SPLIT_SIZE - \	981	uint32_t cmpMixsizeCur = K_L0_SPLIT_SIZE - \
977	info.actualSingleProcessSInnerOriSize % K_L0_SPLIT_SIZE;	982	info.actualSingleProcessSInnerOriSize % K_L0_SPLIT_SIZE;
978	uint32_t oriOnlyLoopTimes = info.actualSingleProcessSInnerOriSize / K_L0_SPLIT_SIZE;	983	uint32_t oriOnlyLoopTimes = info.actualSingleProcessSInnerOriSize / K_L0_SPLIT_SIZE;
	@@ -1061,7 +1066,8 @@ __aicore__ inline void SWACubeBlock<SMLAT>::ComputeMm2(const RunInfo &info, cons
1061		1066
1062	uint32_t headStride = constInfo.headDim;	1067	uint32_t headStride = constInfo.headDim;
1063	uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim;	1068	uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim;
1064	- uint32_t batchStride = constInfo.kvSeqSize / constInfo.cmpRatio * seqStride;	1069	+ uint64_t batchStride = (constInfo.cmpKvStride0 == 0) ?
		1070	+ static_cast<uint64_t>(constInfo.cmpSeqSize) * seqStride : constInfo.cmpKvStride0;
1065		1071
1066	uint64_t curS2 = (uint64_t)info.relativeS2Idx * constInfo.s2BaseSize + info.s2StartPoint + \	1072	uint64_t curS2 = (uint64_t)info.relativeS2Idx * constInfo.s2BaseSize + info.s2StartPoint + \
1067	K_L0_SPLIT_SIZE * kL1;	1073	K_L0_SPLIT_SIZE * kL1;
	@@ -1228,4 +1234,4 @@ __aicore__ inline void SWACubeBlock<SMLAT>::ComputeMm2(const RunInfo &info, cons
1228	qpL1BufIter += mL1Loops;	1234	qpL1BufIter += mL1Loops;
1229	}	1235	}
1230	} // namespace SMLAKernel	1236	} // namespace SMLAKernel
1231	-#endif	1237	+#endif

Mattention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_swa_block_vector.h+22-11

	@@ -424,17 +424,26 @@ __aicore__ inline void SWAVectorBlock<SMLAT>::SetInfInBlk(const LocalTensor<T> &
424	if (startId > endId) {	424	if (startId > endId) {
425	return;	425	return;
426	}	426	}
427	- startId = static_cast<uint64_t>(startId);	427	+ int64_t start = startId < 0 ? 0 : startId;
428	- endId = static_cast<uint64_t>(endId);	428	+ int64_t end = endId >= static_cast<int64_t>(columnCount) ? static_cast<int64_t>(columnCount) - 1 : endId;
429	- uint64_t startFloorAlignSize = startId / BLOCK_ELEMENT_NUM * BLOCK_ELEMENT_NUM;	429	+ if (start > end) {
430	- uint64_t notComputePreMaskOneBlk = (1llu << static_cast<uint64_t>(startId - startFloorAlignSize)) - 1;	430	+ return;
431	- uint64_t notComputePostMaskOneBlk = ~((1llu << static_cast<uint64_t>(endId - startFloorAlignSize + 1)) - 1);	431	+ }
432	- uint64_t notComputeMaskOneBlk = notComputePreMaskOneBlk ^ notComputePostMaskOneBlk;
433		432
434	- uint64_t maskOneBlk = ~notComputeMaskOneBlk;	433	+ uint64_t curStart = static_cast<uint64_t>(start);
435	- uint64_t mask[1] = {maskOneBlk};	434	+ uint64_t curEnd = static_cast<uint64_t>(end);
436	- Duplicate(mmResUb[startFloorAlignSize], SOFTMAX_MIN_NUM, mask,	435	+ while (curStart <= curEnd) {
437	- dealRowCount, 1, columnCount / 8);	436	+ uint64_t blockStart = curStart / BLOCK_ELEMENT_NUM * BLOCK_ELEMENT_NUM;
		437	+ uint64_t blockEnd = blockStart + BLOCK_ELEMENT_NUM - 1;
		438	+ blockEnd = blockEnd > curEnd ? curEnd : blockEnd;
		439	+
		440	+ uint64_t preMask = (1llu << (curStart - blockStart)) - 1;
		441	+ uint64_t postMask = ~((1llu << (blockEnd - blockStart + 1)) - 1);
		442	+ uint64_t mask[1] = {~(preMask \| postMask)};
		443	+ Duplicate(mmResUb[blockStart], SOFTMAX_MIN_NUM, mask,
		444	+ dealRowCount, 1, columnCount / BLOCK_ELEMENT_NUM);
		445	+ curStart = blockEnd + 1;
		446	+ }
438	}	447	}
439		448
440	template <typename SMLAT>	449	template <typename SMLAT>
	@@ -568,7 +577,9 @@ __aicore__ inline void SWAVectorBlock<SMLAT>::ProcessVec1SingleBuf(const RunInfo
568	uint32_t loopCount = (mSplitInfo.vecDealM + mSplitSize - 1) / mSplitSize;	577	uint32_t loopCount = (mSplitInfo.vecDealM + mSplitSize - 1) / mSplitSize;
569	uint32_t tailSplitSize = mSplitInfo.vecDealM - (loopCount - 1) * mSplitSize;	578	uint32_t tailSplitSize = mSplitInfo.vecDealM - (loopCount - 1) * mSplitSize;
570		579
571	- SliceAndContactSinksValue((mSplitInfo.nBufferStartM + mSplitInfo.vecStartM) % constInfo.qHeadNum,	580	+ uint32_t sinkHeadIdx = (info.n2IdxReal * constInfo.gSize + mSplitInfo.nBufferStartM + mSplitInfo.vecStartM) %
		581	+ constInfo.qHeadNum;
		582	+ SliceAndContactSinksValue(sinkHeadIdx,
572	mSplitInfo.vecDealM);	583	mSplitInfo.vecDealM);
573		584
574	for (uint32_t i = 0, dealSize = mSplitSize; i < loopCount; i++) {	585	for (uint32_t i = 0, dealSize = mSplitSize; i < loopCount; i++) {

Mattention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_swa_kernel.h+116-27

	@@ -85,7 +85,8 @@ public:
85	__gm__ uint8_t cmpSparseIndices, __gm__ uint8_t oriBlockTable,	85	__gm__ uint8_t cmpSparseIndices, __gm__ uint8_t oriBlockTable,
86	__gm__ uint8_t cmpBlockTable, __gm__ uint8_t cuSeqlensQ,	86	__gm__ uint8_t cmpBlockTable, __gm__ uint8_t cuSeqlensQ,
87	__gm__ uint8_t* cuSeqlensKV, __gm__ uint8_t cuSeqlensCmpKV, __gm__ uint8_t seqUsedQ,	87	__gm__ uint8_t* cuSeqlensKV, __gm__ uint8_t cuSeqlensCmpKV, __gm__ uint8_t seqUsedQ,
88	- __gm__ uint8_t seqUsedKV, __gm__ uint8_t sinks, __gm__ uint8_t *metadata,	88	+ __gm__ uint8_t seqUsedKV, __gm__ uint8_t seqUsedCmpKV,
		89	+ __gm__ uint8_t cmpResidualKV, __gm__ uint8_t sinks, __gm__ uint8_t *metadata,
89	__gm__ uint8_t attentionOut, __gm__ uint8_t softmaxLse, __gm__ uint8_t *workspace,	90	__gm__ uint8_t attentionOut, __gm__ uint8_t softmaxLse, __gm__ uint8_t *workspace,
90	const SparseFlashMlaTilingData __restrict tiling, __gm__ uint8_t gmTiling,	91	const SparseFlashMlaTilingData __restrict tiling, __gm__ uint8_t gmTiling,
91	TPipe *tPipe);	92	TPipe *tPipe);
	@@ -112,7 +113,6 @@ private:
112	static constexpr uint64_t SYNC_MM2RES_BUF2_FLAG = 11;	113	static constexpr uint64_t SYNC_MM2RES_BUF2_FLAG = 11;
113	static constexpr uint64_t SYNC_FDOUTPUT_BUF_FLAG = 12;	114	static constexpr uint64_t SYNC_FDOUTPUT_BUF_FLAG = 12;
114		115
115	- static constexpr uint64_t kvHeadNum = 1ULL;
116	static constexpr uint64_t headDim = 512ULL;	116	static constexpr uint64_t headDim = 512ULL;
117	static constexpr uint64_t headDimAlign = 512ULL;	117	static constexpr uint64_t headDimAlign = 512ULL;
118	static constexpr uint32_t msdIterNum = 2U;	118	static constexpr uint32_t msdIterNum = 2U;
	@@ -155,6 +155,7 @@ private:
155	GlobalTensor<int32_t> actualSeqLengthsQGm;	155	GlobalTensor<int32_t> actualSeqLengthsQGm;
156	GlobalTensor<int32_t> actualSeqLengthsKVGm;	156	GlobalTensor<int32_t> actualSeqLengthsKVGm;
157	GlobalTensor<int32_t> actualSeqLengthsCmpKVGm;	157	GlobalTensor<int32_t> actualSeqLengthsCmpKVGm;
		158	+ GlobalTensor<int32_t> cmpResidualKVGm;
158		159
159	// workspace	160	// workspace
160	GlobalTensor<MM1_OUT_T> mm1ResGm;	161	GlobalTensor<MM1_OUT_T> mm1ResGm;
	@@ -183,6 +184,8 @@ private:
183	RunInfo &info);	184	RunInfo &info);
184	__aicore__ inline int32_t GetActualSeqLenQ(uint32_t bIdx);	185	__aicore__ inline int32_t GetActualSeqLenQ(uint32_t bIdx);
185	__aicore__ inline int32_t GetActualSeqLenKV(uint32_t bIdx);	186	__aicore__ inline int32_t GetActualSeqLenKV(uint32_t bIdx);
		187	+ __aicore__ inline int32_t GetActualSeqLenCmpKV(uint32_t bIdx, int32_t actualOriS2Size);
		188	+ __aicore__ inline int32_t GetCmpMaskS2Size(uint32_t bIdx, int32_t actualOriS2Size, int32_t actualCmpS2Size);
186	__aicore__ inline void GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx);	189	__aicore__ inline void GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx);
187	// ================================Mm1==============================================	190	// ================================Mm1==============================================
188	__aicore__ inline void ComputeMm1(const RunInfo &info);	191	__aicore__ inline void ComputeMm1(const RunInfo &info);
	@@ -201,7 +204,9 @@ __aicore__ inline void SparseFlashMlaSwa<SMLAT>::InitTilingData()
201	constInfo.bmm2ResUbSize = tilingData->baseParams.bmm2ResUbSize;	204	constInfo.bmm2ResUbSize = tilingData->baseParams.bmm2ResUbSize;
202	// baseParams	205	// baseParams
203	constInfo.batchSize = tilingData->baseParams.batchSize;	206	constInfo.batchSize = tilingData->baseParams.batchSize;
204	- constInfo.qHeadNum = constInfo.gSize = tilingData->baseParams.nNumOfQInOneGroup;	207	+ constInfo.gSize = tilingData->baseParams.nNumOfQInOneGroup;
		208	+ constInfo.kvHeadNum = (tilingData->baseParams.kvHeadNum == 0) ? 1 : tilingData->baseParams.kvHeadNum;
		209	+ constInfo.qHeadNum = constInfo.gSize * constInfo.kvHeadNum;
205	constInfo.kvSeqSize = tilingData->baseParams.kvSeqSize;	210	constInfo.kvSeqSize = tilingData->baseParams.kvSeqSize;
206	constInfo.qSeqSize = tilingData->baseParams.qSeqSize;	211	constInfo.qSeqSize = tilingData->baseParams.qSeqSize;
207	constInfo.oriMaxBlockNumPerBatch = tilingData->baseParams.oriMaxBlockNumPerBatch;	212	constInfo.oriMaxBlockNumPerBatch = tilingData->baseParams.oriMaxBlockNumPerBatch;
	@@ -210,7 +215,6 @@ __aicore__ inline void SparseFlashMlaSwa<SMLAT>::InitTilingData()
210	constInfo.paOriBlockSize = tilingData->baseParams.oriBlockSize;	215	constInfo.paOriBlockSize = tilingData->baseParams.oriBlockSize;
211	constInfo.paCmpBlockSize = tilingData->baseParams.cmpBlockSize;	216	constInfo.paCmpBlockSize = tilingData->baseParams.cmpBlockSize;
212	constInfo.outputLayout = static_cast<SMLA_LAYOUT>(tilingData->baseParams.outputLayout);	217	constInfo.outputLayout = static_cast<SMLA_LAYOUT>(tilingData->baseParams.outputLayout);
213	- constInfo.kvHeadNum = kvHeadNum;
214	constInfo.headDim = headDim;	218	constInfo.headDim = headDim;
215	constInfo.oriMaskMode = tilingData->baseParams.oriMaskMode;	219	constInfo.oriMaskMode = tilingData->baseParams.oriMaskMode;
216	constInfo.oriKvStride0 = tilingData->baseParams.oriKvStride0;	220	constInfo.oriKvStride0 = tilingData->baseParams.oriKvStride0;
	@@ -220,6 +224,8 @@ __aicore__ inline void SparseFlashMlaSwa<SMLAT>::InitTilingData()
220		224
221	constInfo.actualLenDimsQ = tilingData->baseParams.actualLenDimsQ;	225	constInfo.actualLenDimsQ = tilingData->baseParams.actualLenDimsQ;
222	constInfo.actualLenDimsKV = tilingData->baseParams.actualLenDimsKV;	226	constInfo.actualLenDimsKV = tilingData->baseParams.actualLenDimsKV;
		227	+ constInfo.actualLenDimsCmpKV = tilingData->baseParams.actualLenDimsCmpKV;
		228	+ constInfo.cmpResidualKVSize = tilingData->baseParams.cmpResidualKVSize;
223		229
224	// innerSplitParams	230	// innerSplitParams
225	constInfo.mBaseSize = tilingData->baseParams.mBaseSize;	231	constInfo.mBaseSize = tilingData->baseParams.mBaseSize;
	@@ -239,6 +245,7 @@ __aicore__ inline void SparseFlashMlaSwa<SMLAT>::InitTilingData()
239	constInfo.cmpMaskMode = tilingData->cmpParams.cmpMaskMode;	245	constInfo.cmpMaskMode = tilingData->cmpParams.cmpMaskMode;
240	constInfo.cmpKvStride0 = tilingData->cmpParams.cmpKvStride0;	246	constInfo.cmpKvStride0 = tilingData->cmpParams.cmpKvStride0;
241	constInfo.cmpMaxBlockNumPerBatch = tilingData->cmpParams.cmpMaxBlockNumPerBatch;	247	constInfo.cmpMaxBlockNumPerBatch = tilingData->cmpParams.cmpMaxBlockNumPerBatch;
		248	+ constInfo.cmpSeqSize = tilingData->cmpParams.cmpKvSeqSize;
242	}	249	}
243	}	250	}
244		251
	@@ -271,8 +278,11 @@ SparseFlashMlaSwa<SMLAT>::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __
271	{	278	{
272	if (constInfo.actualLenDimsKV != 0) {	279	if (constInfo.actualLenDimsKV != 0) {
273	actualSeqLengthsKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsKV, constInfo.actualLenDimsKV);	280	actualSeqLengthsKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsKV, constInfo.actualLenDimsKV);
274	- if constexpr (TEMPLATE_MODE == CFA_TEMPLATE) {	281	+ }
275	- actualSeqLengthsCmpKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsCmpKV, constInfo.actualLenDimsKV);	282	+ if constexpr (TEMPLATE_MODE == CFA_TEMPLATE) {
		283	+ if (constInfo.actualLenDimsCmpKV != 0) {
		284	+ actualSeqLengthsCmpKVGm.SetGlobalBuffer(
		285	+ (__gm__ int32_t *)actualSeqLengthsCmpKV, constInfo.actualLenDimsCmpKV);
276	}	286	}
277	}	287	}
278	if (constInfo.actualLenDimsQ != 0) {	288	if (constInfo.actualLenDimsQ != 0) {
	@@ -289,24 +299,47 @@ __aicore__ inline void SparseFlashMlaSwa<SMLAT>::InitAllZeroOutput(uint32_t bIdx
289	return;	299	return;
290	}	300	}
291	uint32_t tBase = actualSeqLengthsQGm.GetValue(bIdx);	301	uint32_t tBase = actualSeqLengthsQGm.GetValue(bIdx);
292	- uint64_t attenOutOffset = (tBase + inValidRowS1StartIdx) * kvHeadNum * constInfo.gSize * headDim + // T轴、s1轴偏移	302	+ uint64_t attenOutOffset = (tBase + inValidRowS1StartIdx) * constInfo.kvHeadNum * constInfo.gSize *
293	- n2Idx * constInfo.gSize * headDim; // N2轴偏移	303	+ constInfo.headDim +
		304	+ n2Idx * constInfo.gSize * constInfo.headDim; // N2轴偏移
294	uint64_t lseOffset = (tBase + inValidRowS1StartIdx) * constInfo.gSize + // T轴、s1轴偏移	305	uint64_t lseOffset = (tBase + inValidRowS1StartIdx) * constInfo.gSize + // T轴、s1轴偏移
295	n2Idx * constInfo.qSeqSize * constInfo.gSize; // N2轴偏移	306	n2Idx * constInfo.qSeqSize * constInfo.gSize; // N2轴偏移
296	- matmul::InitOutput<OUT_T>(attentionOutGm[attenOutOffset], inValidRowCount * constInfo.gSize * headDim, 0);	307	+ if (constInfo.kvHeadNum == 1 \|\| inValidRowCount <= 1) {
		308	+ matmul::InitOutput<OUT_T>(
		309	+ attentionOutGm[attenOutOffset], inValidRowCount * constInfo.gSize * constInfo.headDim, 0);
		310	+ } else {
		311	+ uint64_t attenOutRowStride = constInfo.qHeadNum * constInfo.headDim;
		312	+ for (int32_t rowIdx = 0; rowIdx < inValidRowCount; ++rowIdx) {
		313	+ matmul::InitOutput<OUT_T>(
		314	+ attentionOutGm[attenOutOffset + static_cast<uint64_t>(rowIdx) * attenOutRowStride],
		315	+ constInfo.gSize * constInfo.headDim, 0);
		316	+ }
		317	+ }
297	if (constInfo.returnSoftmaxLse) {	318	if (constInfo.returnSoftmaxLse) {
298	- matmul::InitOutput<T>(softmaxLseGm[lseOffset], constInfo.gSize, 0);	319	+ matmul::InitOutput<T>(softmaxLseGm[lseOffset], inValidRowCount * constInfo.gSize, 0);
299	}	320	}
300	} else if (constInfo.outputLayout == SMLA_LAYOUT::BSND) {	321	} else if (constInfo.outputLayout == SMLA_LAYOUT::BSND) {
301	- uint64_t attenOutOffset = bIdx * constInfo.qSeqSize * kvHeadNum * constInfo.gSize * headDim +	322	+ uint64_t attenOutOffset = bIdx * constInfo.qSeqSize * constInfo.kvHeadNum * constInfo.gSize *
302	- inValidRowS1StartIdx * kvHeadNum * constInfo.gSize * headDim + // B轴、S1轴偏移	323	+ constInfo.headDim +
303	- n2Idx * constInfo.gSize * headDim; // N2轴偏移	324	+ inValidRowS1StartIdx * constInfo.kvHeadNum * constInfo.gSize *
		325	+ constInfo.headDim +
		326	+ n2Idx * constInfo.gSize * constInfo.headDim; // N2轴偏移
304	uint64_t lseOffset = bIdx * constInfo.qSeqSize * constInfo.kvHeadNum * constInfo.gSize + // B轴偏移	327	uint64_t lseOffset = bIdx * constInfo.qSeqSize * constInfo.kvHeadNum * constInfo.gSize + // B轴偏移
305	n2Idx * constInfo.qSeqSize * constInfo.gSize + // N2轴偏移	328	n2Idx * constInfo.qSeqSize * constInfo.gSize + // N2轴偏移
306	inValidRowS1StartIdx * constInfo.gSize; // S1轴偏移	329	inValidRowS1StartIdx * constInfo.gSize; // S1轴偏移
307	- matmul::InitOutput<OUT_T>(attentionOutGm[attenOutOffset], inValidRowCount * constInfo.gSize * headDim, 0);	330	+ if (constInfo.kvHeadNum == 1 \|\| inValidRowCount <= 1) {
		331	+ matmul::InitOutput<OUT_T>(
		332	+ attentionOutGm[attenOutOffset], inValidRowCount * constInfo.gSize * constInfo.headDim, 0);
		333	+ } else {
		334	+ uint64_t attenOutRowStride = constInfo.qHeadNum * constInfo.headDim;
		335	+ for (int32_t rowIdx = 0; rowIdx < inValidRowCount; ++rowIdx) {
		336	+ matmul::InitOutput<OUT_T>(
		337	+ attentionOutGm[attenOutOffset + static_cast<uint64_t>(rowIdx) * attenOutRowStride],
		338	+ constInfo.gSize * constInfo.headDim, 0);
		339	+ }
		340	+ }
308	if (constInfo.returnSoftmaxLse) {	341	if (constInfo.returnSoftmaxLse) {
309	- matmul::InitOutput<T>(softmaxLseGm[lseOffset], constInfo.gSize, 0);	342	+ matmul::InitOutput<T>(softmaxLseGm[lseOffset], inValidRowCount * constInfo.gSize, 0);
310	}	343	}
311	}	344	}
312	}	345	}
	@@ -355,18 +388,62 @@ __aicore__ inline int32_t SparseFlashMlaSwa<SMLAT>::GetActualSeqLenKV(uint32_t b
355	}	388	}
356	return actualSeqLengthsKVGm.GetValue(bIdx);	389	return actualSeqLengthsKVGm.GetValue(bIdx);
357	} else if constexpr(KV_LAYOUT_T == SMLA_LAYOUT::BSND) {	390	} else if constexpr(KV_LAYOUT_T == SMLA_LAYOUT::BSND) {
		391	+ tempLoopInfo.actualSeqKVPrefixSum = static_cast<uint64_t>(bIdx * constInfo.kvSeqSize);
		392	+ if (constInfo.actualLenDimsKV != 0) {
		393	+ return actualSeqLengthsKVGm.GetValue(bIdx);
		394	+ }
358	return static_cast<int32_t>(constInfo.kvSeqSize);	395	return static_cast<int32_t>(constInfo.kvSeqSize);
359	} else if constexpr(KV_LAYOUT_T == SMLA_LAYOUT::TND) {	396	} else if constexpr(KV_LAYOUT_T == SMLA_LAYOUT::TND) {
360	int32_t actualSeqKVPrefixSum = actualSeqLengthsKVGm.GetValue(bIdx);	397	int32_t actualSeqKVPrefixSum = actualSeqLengthsKVGm.GetValue(bIdx);
361	int32_t actualSeqKVNextSum = actualSeqLengthsKVGm.GetValue(bIdx + 1);	398	int32_t actualSeqKVNextSum = actualSeqLengthsKVGm.GetValue(bIdx + 1);
362	- if constexpr (TEMPLATE_MODE == CFA_TEMPLATE) {
363	- tempLoopInfo.actualSeqCmpKVPrefixSum = actualSeqLengthsCmpKVGm.GetValue(bIdx);
364	- }
365	tempLoopInfo.actualSeqKVPrefixSum = actualSeqKVPrefixSum;	399	tempLoopInfo.actualSeqKVPrefixSum = actualSeqKVPrefixSum;
366	return actualSeqKVNextSum - actualSeqKVPrefixSum;	400	return actualSeqKVNextSum - actualSeqKVPrefixSum;
367	}	401	}
368	}	402	}
369		403
		404	+template <typename SMLAT>
		405	+__aicore__ inline int32_t SparseFlashMlaSwa<SMLAT>::GetActualSeqLenCmpKV(uint32_t bIdx, int32_t actualOriS2Size)
		406	+{
		407	+ (void)actualOriS2Size;
		408	+ if constexpr (TEMPLATE_MODE != CFA_TEMPLATE) {
		409	+ return 0;
		410	+ }
		411	+ if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::TND) {
		412	+ int32_t actualSeqCmpKVPrefixSum = actualSeqLengthsCmpKVGm.GetValue(bIdx);
		413	+ int32_t actualSeqCmpKVNextSum = actualSeqLengthsCmpKVGm.GetValue(bIdx + 1);
		414	+ tempLoopInfo.actualSeqCmpKVPrefixSum = actualSeqCmpKVPrefixSum;
		415	+ return actualSeqCmpKVNextSum - actualSeqCmpKVPrefixSum;
		416	+ } else if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND) {
		417	+ tempLoopInfo.actualSeqCmpKVPrefixSum = static_cast<uint64_t>(bIdx * constInfo.cmpSeqSize);
		418	+ if (constInfo.actualLenDimsCmpKV != 0) {
		419	+ return actualSeqLengthsCmpKVGm.GetValue(bIdx);
		420	+ }
		421	+ return static_cast<int32_t>(constInfo.cmpSeqSize);
		422	+ } else {
		423	+ tempLoopInfo.actualSeqCmpKVPrefixSum = static_cast<uint64_t>(bIdx * constInfo.cmpSeqSize);
		424	+ if (constInfo.actualLenDimsCmpKV != 0) {
		425	+ return actualSeqLengthsCmpKVGm.GetValue(bIdx);
		426	+ }
		427	+ return (constInfo.cmpSeqSize != 0) ? static_cast<int32_t>(constInfo.cmpSeqSize) :
		428	+ actualOriS2Size / static_cast<int32_t>(constInfo.cmpRatio);
		429	+ }
		430	+}
		431	+
		432	+template <typename SMLAT>
		433	+__aicore__ inline int32_t SparseFlashMlaSwa<SMLAT>::GetCmpMaskS2Size(uint32_t bIdx, int32_t actualOriS2Size,
		434	+ int32_t actualCmpS2Size)
		435	+{
		436	+ (void)actualOriS2Size;
		437	+ if constexpr (TEMPLATE_MODE != CFA_TEMPLATE) {
		438	+ return actualOriS2Size;
		439	+ }
		440	+ int32_t residual = 0;
		441	+ if (constInfo.cmpResidualKVSize != 0) {
		442	+ residual = cmpResidualKVGm.GetValue(bIdx);
		443	+ }
		444	+ return actualCmpS2Size * static_cast<int32_t>(constInfo.cmpRatio) + residual;
		445	+}
		446	+
370	template <typename SMLAT>	447	template <typename SMLAT>
371	__aicore__ inline void SparseFlashMlaSwa<SMLAT>::GetSparseActualSeqLen()	448	__aicore__ inline void SparseFlashMlaSwa<SMLAT>::GetSparseActualSeqLen()
372	{	449	{
	@@ -380,7 +457,7 @@ __aicore__ inline void SparseFlashMlaSwa<SMLAT>::GetSparseActualSeqLen()
380	// 对于cmp部分还有top k, tempLoopInfo.actS2Size只针对cmp	457	// 对于cmp部分还有top k, tempLoopInfo.actS2Size只针对cmp
381	if constexpr (TEMPLATE_MODE == CFA_TEMPLATE) {	458	if constexpr (TEMPLATE_MODE == CFA_TEMPLATE) {
382	int32_t thresHold = (tempLoopInfo.cmpMaskRight + tempLoopInfo.s1EndIdx + 1) / constInfo.cmpRatio;	459	int32_t thresHold = (tempLoopInfo.cmpMaskRight + tempLoopInfo.s1EndIdx + 1) / constInfo.cmpRatio;
383	- tempLoopInfo.actCmpS2Size = thresHold;	460	+ tempLoopInfo.actCmpS2Size = Min(tempLoopInfo.actCmpS2Size, Max(thresHold, 0));
384	}	461	}
385	}	462	}
386		463
	@@ -403,7 +480,8 @@ __aicore__ inline void SparseFlashMlaSwa<SMLAT>::Init(
403	__gm__ uint8_t query, __gm__ uint8_t oriKV, __gm__ uint8_t cmpKV, __gm__ uint8_t cmpSparseIndices,	480	__gm__ uint8_t query, __gm__ uint8_t oriKV, __gm__ uint8_t cmpKV, __gm__ uint8_t cmpSparseIndices,
404	__gm__ uint8_t oriBlockTable, __gm__ uint8_t cmpBlockTable, __gm__ uint8_t *cuSeqlensQ,	481	__gm__ uint8_t oriBlockTable, __gm__ uint8_t cmpBlockTable, __gm__ uint8_t *cuSeqlensQ,
405	__gm__ uint8_t cuSeqlensKV, __gm__ uint8_t cuSeqlensCmpKV, __gm__ uint8_t *seqUsedQ,	482	__gm__ uint8_t cuSeqlensKV, __gm__ uint8_t cuSeqlensCmpKV, __gm__ uint8_t *seqUsedQ,
406	- __gm__ uint8_t seqUsedKV, __gm__ uint8_t sinks, __gm__ uint8_t metadata, __gm__ uint8_t attentionOut, __gm__ uint8_t *softmaxLse,	483	+ __gm__ uint8_t seqUsedKV, __gm__ uint8_t seqUsedCmpKV, __gm__ uint8_t *cmpResidualKV,
		484	+ __gm__ uint8_t sinks, __gm__ uint8_t metadata, __gm__ uint8_t attentionOut, __gm__ uint8_t softmaxLse,
407	__gm__ uint8_t workspace, const SparseFlashMlaTilingData __restrict tiling, __gm__ uint8_t *gmTiling,	485	__gm__ uint8_t workspace, const SparseFlashMlaTilingData __restrict tiling, __gm__ uint8_t *gmTiling,
408	TPipe *tPipe)	486	TPipe *tPipe)
409	{	487	{
	@@ -425,9 +503,9 @@ __aicore__ inline void SparseFlashMlaSwa<SMLAT>::Init(
425	InitActualSeqLen(seqUsedQ, cuSeqlensKV, cuSeqlensCmpKV);	503	InitActualSeqLen(seqUsedQ, cuSeqlensKV, cuSeqlensCmpKV);
426	} else if ((KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND \|\| KV_LAYOUT_T == SMLA_LAYOUT::BSND)	504	} else if ((KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND \|\| KV_LAYOUT_T == SMLA_LAYOUT::BSND)
427	&& LAYOUT_T == SMLA_LAYOUT::TND) {	505	&& LAYOUT_T == SMLA_LAYOUT::TND) {
428	- InitActualSeqLen(cuSeqlensQ, seqUsedKV);	506	+ InitActualSeqLen(cuSeqlensQ, seqUsedKV, seqUsedCmpKV);
429	} else if ((KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND \|\| KV_LAYOUT_T == SMLA_LAYOUT::BSND)) {	507	} else if ((KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND \|\| KV_LAYOUT_T == SMLA_LAYOUT::BSND)) {
430	- InitActualSeqLen(seqUsedQ, seqUsedKV);	508	+ InitActualSeqLen(seqUsedQ, seqUsedKV, seqUsedCmpKV);
431	}	509	}
432	metadataGm.SetGlobalBuffer((__gm__ uint32_t *)metadata);	510	metadataGm.SetGlobalBuffer((__gm__ uint32_t *)metadata);
433	InitCalcParamsEach();	511	InitCalcParamsEach();
	@@ -438,6 +516,9 @@ __aicore__ inline void SparseFlashMlaSwa<SMLAT>::Init(
438	oriKvGm.SetGlobalBuffer((__gm__ KV_T *)oriKV);	516	oriKvGm.SetGlobalBuffer((__gm__ KV_T *)oriKV);
439	if constexpr (TEMPLATE_MODE == CFA_TEMPLATE) {	517	if constexpr (TEMPLATE_MODE == CFA_TEMPLATE) {
440	cmpKvGm.SetGlobalBuffer((__gm__ KV_T *)cmpKV);	518	cmpKvGm.SetGlobalBuffer((__gm__ KV_T *)cmpKV);
		519	+ if (constInfo.cmpResidualKVSize != 0) {
		520	+ cmpResidualKVGm.SetGlobalBuffer((__gm__ int32_t *)cmpResidualKV, constInfo.cmpResidualKVSize);
		521	+ }
441	}	522	}
442		523
443	if (sinks != nullptr) {	524	if (sinks != nullptr) {
	@@ -562,7 +643,9 @@ __aicore__ inline void SparseFlashMlaSwa<SMLAT>::CalcParams(uint32_t loop, uint3
562	uint64_t tndBIdxOffsetForCmpKV = tempLoopInfo.actualSeqCmpKVPrefixSum * constInfo.kvHeadNum * constInfo.headDim;	643	uint64_t tndBIdxOffsetForCmpKV = tempLoopInfo.actualSeqCmpKVPrefixSum * constInfo.kvHeadNum * constInfo.headDim;
563		644
564	if (info.isFirstSInnerLoop) {	645	if (info.isFirstSInnerLoop) {
565	- tensorACoreOffset = tndBIdxOffsetForQ + info.gS1Idx * constInfo.headDim;	646	+ uint64_t s1HeadOffset = (info.gS1Idx / constInfo.gSize) * constInfo.qHeadNum;
		647	+ uint64_t qHeadOffset = info.n2Idx * constInfo.gSize + info.gS1Idx % constInfo.gSize;
		648	+ tensorACoreOffset = tndBIdxOffsetForQ + (s1HeadOffset + qHeadOffset) * constInfo.headDim;
566	tensorBCoreOffset = tndBIdxOffsetForKV + info.n2Idx * constInfo.headDim;	649	tensorBCoreOffset = tndBIdxOffsetForKV + info.n2Idx * constInfo.headDim;
567	tensorCmpBCoreOffset = tndBIdxOffsetForCmpKV + info.n2Idx * constInfo.headDim;	650	tensorCmpBCoreOffset = tndBIdxOffsetForCmpKV + info.n2Idx * constInfo.headDim;
568	}	651	}
	@@ -691,8 +774,8 @@ __aicore__ inline void SparseFlashMlaSwa<SMLAT>::Process()
691	template <typename SMLAT>	774	template <typename SMLAT>
692	__aicore__ inline void SparseFlashMlaSwa<SMLAT>::GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx)	775	__aicore__ inline void SparseFlashMlaSwa<SMLAT>::GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx)
693	{	776	{
694	- bIdx = bN2Idx / kvHeadNum;	777	+ bIdx = bN2Idx / constInfo.kvHeadNum;
695	- n2Idx = bN2Idx % kvHeadNum;	778	+ n2Idx = bN2Idx % constInfo.kvHeadNum;
696	}	779	}
697		780
698	template <typename SMLAT>	781	template <typename SMLAT>
	@@ -736,6 +819,10 @@ __aicore__ inline void SparseFlashMlaSwa<SMLAT>::ProcessBalance()
736	gS1LoopEnd = isS1S2ZeroAndLastBatch ? gS1LoopEnd + 1 : gS1LoopEnd;	819	gS1LoopEnd = isS1S2ZeroAndLastBatch ? gS1LoopEnd + 1 : gS1LoopEnd;
737	for (uint32_t gS1LoopIdx = constInfo.gS1Start; gS1LoopIdx < gS1LoopEnd; gS1LoopIdx++) {	820	for (uint32_t gS1LoopIdx = constInfo.gS1Start; gS1LoopIdx < gS1LoopEnd; gS1LoopIdx++) {
738	tempLoopInfo.actOriS2Size = GetActualSeqLenKV(tempLoopInfo.bIdx);	821	tempLoopInfo.actOriS2Size = GetActualSeqLenKV(tempLoopInfo.bIdx);
		822	+ if constexpr (TEMPLATE_MODE == CFA_TEMPLATE) {
		823	+ tempLoopInfo.actCmpS2Size =
		824	+ GetActualSeqLenCmpKV(tempLoopInfo.bIdx, tempLoopInfo.actOriS2Size);
		825	+ }
739	// 对于各轴上的真实的idx, 采用左闭右闭的方案	826	// 对于各轴上的真实的idx, 采用左闭右闭的方案
740	// 跳过行无效部分，从有效行开始后续计算	827	// 跳过行无效部分，从有效行开始后续计算
741	tempLoopInfo.gS1Idx = inValidRowCount * constInfo.gSize + gS1LoopIdx * constInfo.mBaseSize;	828	tempLoopInfo.gS1Idx = inValidRowCount * constInfo.gSize + gS1LoopIdx * constInfo.mBaseSize;
	@@ -748,7 +835,9 @@ __aicore__ inline void SparseFlashMlaSwa<SMLAT>::ProcessBalance()
748	tempLoopInfo.oriMaskLeft = Max(tempLoopInfo.actOriS2Size - tempLoopInfo.actS1Size +	835	tempLoopInfo.oriMaskLeft = Max(tempLoopInfo.actOriS2Size - tempLoopInfo.actS1Size +
749	static_cast<int32_t>(tempLoopInfo.s1StartIdx) - constInfo.oriWinLeft, 0);	836	static_cast<int32_t>(tempLoopInfo.s1StartIdx) - constInfo.oriWinLeft, 0);
750	if constexpr (TEMPLATE_MODE == CFA_TEMPLATE) {	837	if constexpr (TEMPLATE_MODE == CFA_TEMPLATE) {
751	- tempLoopInfo.cmpMaskRight = tempLoopInfo.actOriS2Size - tempLoopInfo.actS1Size;	838	+ int32_t cmpMaskS2Size = GetCmpMaskS2Size(
		839	+ tempLoopInfo.bIdx, tempLoopInfo.actOriS2Size, tempLoopInfo.actCmpS2Size);
		840	+ tempLoopInfo.cmpMaskRight = cmpMaskS2Size - tempLoopInfo.actS1Size;
752	}	841	}
753	GetSparseActualSeqLen();	842	GetSparseActualSeqLen();
754	UpdateInnerLoopCond();	843	UpdateInnerLoopCond();
	@@ -835,4 +924,4 @@ __aicore__ inline void SparseFlashMlaSwa<SMLAT>::PreloadPipeline(uint32_t loop,
835	}	924	}
836	}	925	}
837	} // namespace SMLAKernel	926	} // namespace SMLAKernel
838	-#endif // SPARSE_FLASH_MLA_SWA_KERNEL_H	927	+#endif // SPARSE_FLASH_MLA_SWA_KERNEL_H

Mattention/sparse_flash_mla/op_kernel/sparse_flash_mla.cpp+9-7

Mattention/sparse_flash_mla/op_kernel/sparse_flash_mla_common.h+3-2

	@@ -41,7 +41,7 @@ constexpr SoftmaxConfig SMLA_SOFTMAX_FLASHV2_CFG_WITHOUT_BRC = {false, 0, 0, Sof
41		41
42	template <typename Q_T, typename KV_T, typename OUT_T, const bool FLASH_DECODE = false,	42	template <typename Q_T, typename KV_T, typename OUT_T, const bool FLASH_DECODE = false,
43	SMLA_LAYOUT LAYOUT_T = SMLA_LAYOUT::BSND, SMLA_LAYOUT KV_LAYOUT_T = SMLA_LAYOUT::PA_BBND,	43	SMLA_LAYOUT LAYOUT_T = SMLA_LAYOUT::BSND, SMLA_LAYOUT KV_LAYOUT_T = SMLA_LAYOUT::PA_BBND,
44	- int TEMPLATE_MODE = 0, typename... Args>	44	+ int TEMPLATE_MODE = 0, const bool HEAD_RATIO_ONE = false, typename... Args>
45	struct SMLAType {	45	struct SMLAType {
46	using queryType = Q_T;	46	using queryType = Q_T;
47	using kvType = KV_T;	47	using kvType = KV_T;
	@@ -51,6 +51,7 @@ struct SMLAType {
51	static constexpr SMLA_LAYOUT kvLayout = KV_LAYOUT_T;	51	static constexpr SMLA_LAYOUT kvLayout = KV_LAYOUT_T;
52	static constexpr bool pageAttention = (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND);	52	static constexpr bool pageAttention = (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND);
53	static constexpr int templateMode = TEMPLATE_MODE;	53	static constexpr int templateMode = TEMPLATE_MODE;
		54	+ static constexpr bool headRatioOne = HEAD_RATIO_ONE;
54	};	55	};
55		56
56	// ================================Util functions==================================	57	// ================================Util functions==================================
	@@ -323,4 +324,4 @@ struct MSplitInfo {
323	};	324	};
324	#endif	325	#endif
325	} // namespace SMLAKernel	326	} // namespace SMLAKernel
326	-#endif // SPARSE_ATTN_SHAREDKV_COMMON_H	327	+#endif // SPARSE_ATTN_SHAREDKV_COMMON_H

Mattention/sparse_flash_mla/op_kernel/sparse_flash_mla_template_tiling_key.h+3-1

	@@ -37,6 +37,7 @@ ASCENDC_TPL_ARGS_DECL(SparseFlashMla, // 算子OpType
37	ASCENDC_TPL_UINT_DECL(TEMPLATE_MODE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, SWA_TEMPLATE,	37	ASCENDC_TPL_UINT_DECL(TEMPLATE_MODE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, SWA_TEMPLATE,
38	CFA_TEMPLATE, SCFA_TEMPLATE),	38	CFA_TEMPLATE, SCFA_TEMPLATE),
39	ASCENDC_TPL_BOOL_DECL(SPLIT_G, 0, 1),	39	ASCENDC_TPL_BOOL_DECL(SPLIT_G, 0, 1),
		40	+ ASCENDC_TPL_BOOL_DECL(HEAD_RATIO_ONE, 0, 1),
40	);	41	);
41		42
42	// 支持的模板参数组合	43	// 支持的模板参数组合
	@@ -49,7 +50,8 @@ ASCENDC_TPL_SEL(
49	ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, SWA_TEMPLATE,	50	ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, SWA_TEMPLATE,
50	CFA_TEMPLATE, SCFA_TEMPLATE),	51	CFA_TEMPLATE, SCFA_TEMPLATE),
51	ASCENDC_TPL_BOOL_SEL(SPLIT_G, 0, 1),	52	ASCENDC_TPL_BOOL_SEL(SPLIT_G, 0, 1),
		53	+ ASCENDC_TPL_BOOL_SEL(HEAD_RATIO_ONE, 0, 1),
52	)	54	)
53	);	55	);
54		56
55	-#endif // TEMPLATE_TILING_KEY	57	+#endif // TEMPLATE_TILING_KEY

Mattention/sparse_flash_mla/tests/pytest/batch/sparse_flash_mla_process.py+8-7

	@@ -72,12 +72,13 @@ def call_npu(input_data):
72	layout_q = tensor_input['layout_q'] if type(tensor_input['layout_q']) == type('TND') else tensor_input['layout_q'][0]	72	layout_q = tensor_input['layout_q'] if type(tensor_input['layout_q']) == type('TND') else tensor_input['layout_q'][0]
73	layout_kv = tensor_input['layout_kv']	73	layout_kv = tensor_input['layout_kv']
74	max_seqlen_q = metadata_input['max_seqlen_q']	74	max_seqlen_q = metadata_input['max_seqlen_q']
75	- ori_max_s2 = metadata_input['max_seqlen_kv']	75	+ max_seqlen_ori_kv = metadata_input['max_seqlen_ori_kv']
		76	+ max_seqlen_cmp_kv = metadata_input['max_seqlen_cmp_kv']
76	ori_sparse_indices = tensor_input['ori_sparse_indices']	77	ori_sparse_indices = tensor_input['ori_sparse_indices']
77	cmp_sparse_indices = tensor_input['cmp_sparse_indices']	78	cmp_sparse_indices = tensor_input['cmp_sparse_indices']
78	cmp_block_table = tensor_input['cmp_block_table']	79	cmp_block_table = tensor_input['cmp_block_table']
79	- ori_topk_length = metadata_input['ori_topk_length'].unsqueeze(-1).npu() if metadata_input['ori_topk_length'] is not None else None	80	+ ori_topk_length = None
80	- cmp_topk_length = metadata_input['cmp_topk_length'].unsqueeze(-1).npu() if metadata_input['cmp_topk_length'] is not None else None	81	+ cmp_topk_length = None
81	return_softmax_lse = params.get('return_softmax_lse')	82	return_softmax_lse = params.get('return_softmax_lse')
82		83
83	# 将需要上NPU的tensor搬到NPU	84	# 将需要上NPU的tensor搬到NPU
	@@ -105,10 +106,10 @@ def call_npu(input_data):
105	cmp_topk_length=cmp_topk_length,	106	cmp_topk_length=cmp_topk_length,
106	batch_size=B,	107	batch_size=B,
107	max_seqlen_q=max_seqlen_q,	108	max_seqlen_q=max_seqlen_q,
108	- # max_seqlen_ori_kv=max_seqlen_ori_kv,	109	+ max_seqlen_ori_kv=max_seqlen_ori_kv,
109	- # max_seqlen_cmp_kv=max_seqlen_cmp_kv,	110	+ max_seqlen_cmp_kv=max_seqlen_cmp_kv,
110	- # ori_topk=K,	111	+ ori_topk=K if ori_sparse_indices is not None else 0,
111	- cmp_topk=K,	112	+ cmp_topk=K if cmp_sparse_indices is not None else 0,
112	cmp_ratio=cmp_ratio if cmp_ratio is not None else 1,	113	cmp_ratio=cmp_ratio if cmp_ratio is not None else 1,
113	ori_mask_mode=ori_mask_mode,	114	ori_mask_mode=ori_mask_mode,
114	cmp_mask_mode=cmp_mask_mode if cmp_mask_mode is not None else 3,	115	cmp_mask_mode=cmp_mask_mode if cmp_mask_mode is not None else 3,

Mattention/sparse_flash_mla/tests/pytest/sparse_flash_mla_golden.py+9-6

	@@ -694,8 +694,8 @@ def gen_cmp_kv(layout_q, layout_kv, cmp_kv_type, B, S1, S2, T1, T2, T3, N2, D, K
694	data_range_right=DATA_RANGE_RIGHT):	694	data_range_right=DATA_RANGE_RIGHT):
695	if cmp_ratio is None:	695	if cmp_ratio is None:
696	raise ValueError(f"cmp_ratio can't be None")	696	raise ValueError(f"cmp_ratio can't be None")
697	- if cmp_ratio == 0:	697	+ if cmp_ratio < 1 or cmp_ratio > 128:
698	- raise ValueError(f"cmp_ratio can't be 0")	698	+ raise ValueError(f"cmp_ratio should be in range [1, 128], but got {cmp_ratio}")
699		699
700	if layout_kv == "PA_BBND":	700	if layout_kv == "PA_BBND":
701	ori_max_s2 = max(seqused_ori_kv)	701	ori_max_s2 = max(seqused_ori_kv)
	@@ -876,7 +876,9 @@ def gen_data(params, template_mode=None):
876	else:	876	else:
877	raise ValueError(f"layout_kv is not support {layout_kv}")	877	raise ValueError(f"layout_kv is not support {layout_kv}")
878		878
879	- if seqused_ori_kv is not None and cmp_ratio is not None and cmp_ratio != 0:	879	+ if seqused_ori_kv is not None and cmp_ratio is not None:
		880	+ if cmp_ratio < 1:
		881	+ raise ValueError(f"cmp_ratio should be in range [1, 128], but got {cmp_ratio}")
880	seqused_cmp_kv = seqused_ori_kv // cmp_ratio	882	seqused_cmp_kv = seqused_ori_kv // cmp_ratio
881	cmp_residual_kv = seqused_ori_kv % cmp_ratio	883	cmp_residual_kv = seqused_ori_kv % cmp_ratio
882	# 路由到三个算子的逻辑：	884	# 路由到三个算子的逻辑：
	@@ -971,11 +973,12 @@ def gen_data(params, template_mode=None):
971	'seqused_ori_kv': seqused_ori_kv,	973	'seqused_ori_kv': seqused_ori_kv,
972	'seqused_cmp_kv': seqused_cmp_kv,	974	'seqused_cmp_kv': seqused_cmp_kv,
973	'cmp_residual_kv': cmp_residual_kv,	975	'cmp_residual_kv': cmp_residual_kv,
974	- 'ori_topk_length': ori_topk_length,	976	+ 'ori_topk_length': None,
975	- 'cmp_topk_length': cmp_topk_length,	977	+ 'cmp_topk_length': None,
976	'B': B,	978	'B': B,
977	'max_seqlen_q': max_seqlen_q,	979	'max_seqlen_q': max_seqlen_q,
978	- 'max_seqlen_kv': max(seqused_ori_kv) if seqused_ori_kv is not None else (T2 if layout_kv == "TND" else S2),	980	+ 'max_seqlen_ori_kv': max(seqused_ori_kv) if seqused_ori_kv is not None else (T2 if layout_kv == "TND" else S2),
		981	+ 'max_seqlen_cmp_kv': max(seqused_cmp_kv) if seqused_cmp_kv is not None else 0,
979	'K': K,	982	'K': K,
980	'cmp_ratio': cmp_ratio,	983	'cmp_ratio': cmp_ratio,
981	'ori_mask_mode': ori_mask_mode,	984	'ori_mask_mode': ori_mask_mode,

Mattention/sparse_flash_mla/tests/ut/op_host/test_sparse_flash_mla_tiling.cpp+9-27

Mattention/sparse_flash_mla_metadata/op_host/sparse_flash_mla_metadata_check.h+70-10

	@@ -18,6 +18,7 @@
18	#include "opdev/data_type_utils.h"	18	#include "opdev/data_type_utils.h"
19	#include "opdev/tensor_view_utils.h"	19	#include "opdev/tensor_view_utils.h"
20	#include "../../sparse_flash_mla/op_kernel/sparse_flash_mla_metadata.h"	20	#include "../../sparse_flash_mla/op_kernel/sparse_flash_mla_metadata.h"
		21	+#include <cstring>
21		22
22	#ifdef __cplusplus	23	#ifdef __cplusplus
23	extern "C" {	24	extern "C" {
	@@ -39,11 +40,31 @@ inline constexpr int64_t SMLA_CMP_RATIO_UPPER_BOUND = 128;
39	inline constexpr int64_t SMLA_NUM_HEADS_Q_LOWER_BOUND = 1;	40	inline constexpr int64_t SMLA_NUM_HEADS_Q_LOWER_BOUND = 1;
40	inline constexpr int64_t SMLA_NUM_HEADS_Q_UPPER_BOUND = 128;	41	inline constexpr int64_t SMLA_NUM_HEADS_Q_UPPER_BOUND = 128;
41		42
		43	+inline bool IsPowerOfTwoInRangeSmla(int64_t value, int64_t minValue, int64_t maxValue)
		44	+{
		45	+ return value >= minValue && value <= maxValue && ((value & (value - 1)) == 0);
		46	+}
		47	+
		48	+inline bool IsA5Smla(const char *socVersion)
		49	+{
		50	+ return socVersion != nullptr && strstr(socVersion, "Ascend950") != nullptr;
		51	+}
		52	+
42	inline bool IsTensorExistSmla(const aclTensor *tensor)	53	inline bool IsTensorExistSmla(const aclTensor *tensor)
43	{	54	{
44	return (tensor != nullptr) && (tensor->GetViewShape().GetDimNum() > 0) && (tensor->GetViewShape().GetDim(0) > 0);	55	return (tensor != nullptr) && (tensor->GetViewShape().GetDimNum() > 0) && (tensor->GetViewShape().GetDim(0) > 0);
45	}	56	}
46		57
		58	+aclnnStatus CheckReservedOptionalTensorSmla(const aclTensor tensor, const char tensorName)
		59	+{
		60	+ if (!IsTensorExistSmla(tensor)) {
		61	+ return ACLNN_SUCCESS;
		62	+ }
		63	+ OP_LOGE(ACLNN_ERR_PARAM_INVALID,
		64	+ "%s is reserved and does not support non-empty tensor in current version", tensorName);
		65	+ return ACLNN_ERR_PARAM_INVALID;
		66	+}
		67	+
47	int64_t GetDimNumSmla(const aclTensor *tensor)	68	int64_t GetDimNumSmla(const aclTensor *tensor)
48	{	69	{
49	if (tensor == nullptr) {	70	if (tensor == nullptr) {
	@@ -100,6 +121,26 @@ aclnnStatus CheckSingleParamSmla(int64_t batchSize, int64_t maxSeqlenQ, int64_t
100	OP_LOGE(ACLNN_ERR_PARAM_INVALID, "num_heads_kv should only be 1, but got %lld", numHeadsKv);	121	OP_LOGE(ACLNN_ERR_PARAM_INVALID, "num_heads_kv should only be 1, but got %lld", numHeadsKv);
101	return ACLNN_ERR_PARAM_INVALID;	122	return ACLNN_ERR_PARAM_INVALID;
102	}	123	}
		124	+ if (numHeadsQ % numHeadsKv != 0) {
		125	+ OP_LOGE(ACLNN_ERR_PARAM_INVALID,
		126	+ "num_heads_q should be divisible by num_heads_kv, but got %lld and %lld", numHeadsQ, numHeadsKv);
		127	+ return ACLNN_ERR_PARAM_INVALID;
		128	+ }
		129	+ int64_t headRatio = numHeadsQ / numHeadsKv;
		130	+ bool isA5 = IsA5Smla(socVersion);
		131	+ if (isA5) {
		132	+ if (headRatio < SMLA_NUM_HEADS_Q_LOWER_BOUND \|\| headRatio > SMLA_NUM_HEADS_Q_UPPER_BOUND) {
		133	+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "num_heads_q / num_heads_kv should be in [%lld, %lld], but got %lld",
		134	+ SMLA_NUM_HEADS_Q_LOWER_BOUND, SMLA_NUM_HEADS_Q_UPPER_BOUND, headRatio);
		135	+ return ACLNN_ERR_PARAM_INVALID;
		136	+ }
		137	+ } else if (!IsPowerOfTwoInRangeSmla(
		138	+ headRatio, SMLA_NUM_HEADS_Q_LOWER_BOUND, SMLA_NUM_HEADS_Q_UPPER_BOUND)) {
		139	+ OP_LOGE(ACLNN_ERR_PARAM_INVALID,
		140	+ "num_heads_q / num_heads_kv should be power of two in [%lld, %lld], but got %lld",
		141	+ SMLA_NUM_HEADS_Q_LOWER_BOUND, SMLA_NUM_HEADS_Q_UPPER_BOUND, headRatio);
		142	+ return ACLNN_ERR_PARAM_INVALID;
		143	+ }
103	// head_dim: 512	144	// head_dim: 512
104	if (headDim != 512) {	145	if (headDim != 512) {
105	OP_LOGE(ACLNN_ERR_PARAM_INVALID, "head_dim should only be 512, but got %lld", headDim);	146	OP_LOGE(ACLNN_ERR_PARAM_INVALID, "head_dim should only be 512, but got %lld", headDim);
	@@ -111,6 +152,10 @@ aclnnStatus CheckSingleParamSmla(int64_t batchSize, int64_t maxSeqlenQ, int64_t
111	OP_LOGE(ACLNN_ERR_PARAM_INVALID, "When has_ori_kv is true, ori_topk should be >= 0, but got %lld", oriTopk);	152	OP_LOGE(ACLNN_ERR_PARAM_INVALID, "When has_ori_kv is true, ori_topk should be >= 0, but got %lld", oriTopk);
112	return ACLNN_ERR_PARAM_INVALID;	153	return ACLNN_ERR_PARAM_INVALID;
113	}	154	}
		155	+ if (!isA5 && oriTopk != 0) {
		156	+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "ori_topk is reserved and should only be 0, but got %lld", oriTopk);
		157	+ return ACLNN_ERR_PARAM_INVALID;
		158	+ }
114	// ori_mask_mode: 0, 3, or 4	159	// ori_mask_mode: 0, 3, or 4
115	if (oriMaskMode != static_cast<int64_t>(SparseModeSmla::DEFAULT_MASK) &&	160	if (oriMaskMode != static_cast<int64_t>(SparseModeSmla::DEFAULT_MASK) &&
116	oriMaskMode != static_cast<int64_t>(SparseModeSmla::RIGHT_DOWN_CAUSAL) &&	161	oriMaskMode != static_cast<int64_t>(SparseModeSmla::RIGHT_DOWN_CAUSAL) &&
	@@ -138,6 +183,10 @@ aclnnStatus CheckSingleParamSmla(int64_t batchSize, int64_t maxSeqlenQ, int64_t
138	OP_LOGE(ACLNN_ERR_PARAM_INVALID, "When has_cmp_kv is true, cmp_topk should be >= 0, but got %lld", cmpTopk);	183	OP_LOGE(ACLNN_ERR_PARAM_INVALID, "When has_cmp_kv is true, cmp_topk should be >= 0, but got %lld", cmpTopk);
139	return ACLNN_ERR_PARAM_INVALID;	184	return ACLNN_ERR_PARAM_INVALID;
140	}	185	}
		186	+ if (!isA5 && cmpTopk != 0 && cmpTopk != 512 && cmpTopk != 1024) {
		187	+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "cmp_topk should be 0, 512 or 1024, but got %lld", cmpTopk);
		188	+ return ACLNN_ERR_PARAM_INVALID;
		189	+ }
141	// cmp_mask_mode: 0 or 3	190	// cmp_mask_mode: 0 or 3
142	if (cmpMaskMode != static_cast<int64_t>(SparseModeSmla::DEFAULT_MASK) &&	191	if (cmpMaskMode != static_cast<int64_t>(SparseModeSmla::DEFAULT_MASK) &&
143	cmpMaskMode != static_cast<int64_t>(SparseModeSmla::RIGHT_DOWN_CAUSAL)) {	192	cmpMaskMode != static_cast<int64_t>(SparseModeSmla::RIGHT_DOWN_CAUSAL)) {
	@@ -145,10 +194,14 @@ aclnnStatus CheckSingleParamSmla(int64_t batchSize, int64_t maxSeqlenQ, int64_t
145	cmpMaskMode);	194	cmpMaskMode);
146	return ACLNN_ERR_PARAM_INVALID;	195	return ACLNN_ERR_PARAM_INVALID;
147	}	196	}
148	- // cmp_ratio: 1~128	197	+ if (isA5) {
149	- if (cmpRatio < SMLA_CMP_RATIO_LOWER_BOUND \|\| cmpRatio > SMLA_CMP_RATIO_UPPER_BOUND) {	198	+ if (cmpRatio < SMLA_CMP_RATIO_LOWER_BOUND \|\| cmpRatio > SMLA_CMP_RATIO_UPPER_BOUND) {
150	- OP_LOGE(ACLNN_ERR_PARAM_INVALID, "When has_cmp_kv is true, cmp_ratio should be in [%lld, %lld], "	199	+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "cmp_ratio should be in [%lld, %lld], but got %lld",
151	- "but got %lld", SMLA_CMP_RATIO_LOWER_BOUND, SMLA_CMP_RATIO_UPPER_BOUND, cmpRatio);	200	+ SMLA_CMP_RATIO_LOWER_BOUND, SMLA_CMP_RATIO_UPPER_BOUND, cmpRatio);
		201	+ return ACLNN_ERR_PARAM_INVALID;
		202	+ }
		203	+ } else if (cmpRatio != 1 && cmpRatio != 4 && cmpRatio != 128) {
		204	+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "cmp_ratio should be 1, 4 or 128, but got %lld", cmpRatio);
152	return ACLNN_ERR_PARAM_INVALID;	205	return ACLNN_ERR_PARAM_INVALID;
153	}	206	}
154	}	207	}
	@@ -343,10 +396,16 @@ aclnnStatus CheckConsistencySmla(const aclTensor *cuSeqlensQOptional, const aclT
343	const aclTensor sequsedOriKvOptional, const aclTensor sequsedCmpKvOptional,	396	const aclTensor sequsedOriKvOptional, const aclTensor sequsedCmpKvOptional,
344	const aclTensor cmpResidualKvOptional, const aclTensor oriTopkLengthOptional,	397	const aclTensor cmpResidualKvOptional, const aclTensor oriTopkLengthOptional,
345	const aclTensor cmpTopkLengthOptional, int64_t batchSize, const char layoutQOptional,	398	const aclTensor cmpTopkLengthOptional, int64_t batchSize, const char layoutQOptional,
346	- const char layoutKvOptional, bool hasOriKv, bool hasCmpKv, const aclTensor metadata)	399	+ const char *layoutKvOptional, bool hasOriKv, bool hasCmpKv, bool isA5,
		400	+ const aclTensor *metadata)
347	{	401	{
348	aclDataType dataType = aclDataType::ACL_DT_UNDEFINED;	402	aclDataType dataType = aclDataType::ACL_DT_UNDEFINED;
349	int64_t dimNum = -1;	403	int64_t dimNum = -1;
		404	+ if (!isA5 &&
		405	+ (CheckReservedOptionalTensorSmla(oriTopkLengthOptional, "ori_topk_length") != ACLNN_SUCCESS \|\|
		406	+ CheckReservedOptionalTensorSmla(cmpTopkLengthOptional, "cmp_topk_length") != ACLNN_SUCCESS)) {
		407	+ return ACLNN_ERR_PARAM_INVALID;
		408	+ }
350	// 校验 cu_seqlens_q	409	// 校验 cu_seqlens_q
351	if (IsTensorExistSmla(cuSeqlensQOptional)) {	410	if (IsTensorExistSmla(cuSeqlensQOptional)) {
352	// 校验 cu_seqlens_q 维度	411	// 校验 cu_seqlens_q 维度
	@@ -535,8 +594,8 @@ aclnnStatus CheckConsistencySmla(const aclTensor *cuSeqlensQOptional, const aclT
535	// 校验 q/kv 维度一致性	594	// 校验 q/kv 维度一致性
536	int64_t queryBatchSize = GetQueryBatchSizeSmla(sequsedQOptional, cuSeqlensQOptional, layoutQOptional, batchSize);	595	int64_t queryBatchSize = GetQueryBatchSizeSmla(sequsedQOptional, cuSeqlensQOptional, layoutQOptional, batchSize);
537	if (hasOriKv) {	596	if (hasOriKv) {
538	- int64_t oriKvBatchSize = GetOriKvBatchSizeSmla(sequsedOriKvOptional, cuSeqlensOriKvOptional, layoutKvOptional,	597	+ int64_t oriKvBatchSize = GetOriKvBatchSizeSmla(
539	- batchSize);	598	+ sequsedOriKvOptional, cuSeqlensOriKvOptional, layoutKvOptional, batchSize);
540	if (queryBatchSize != oriKvBatchSize) {	599	if (queryBatchSize != oriKvBatchSize) {
541	OP_LOGE(ACLNN_ERR_PARAM_INVALID, "when has_ori_kv is true, the batch_size obtained from q should be "	600	OP_LOGE(ACLNN_ERR_PARAM_INVALID, "when has_ori_kv is true, the batch_size obtained from q should be "
542	"the same as that obtained from ori_kv, but got %lld and %lld", queryBatchSize,	601	"the same as that obtained from ori_kv, but got %lld and %lld", queryBatchSize,
	@@ -545,8 +604,8 @@ aclnnStatus CheckConsistencySmla(const aclTensor *cuSeqlensQOptional, const aclT
545	}	604	}
546	}	605	}
547	if (hasCmpKv) {	606	if (hasCmpKv) {
548	- int64_t cmpKvBatchSize = GetCmpKvBatchSizeSmla(sequsedCmpKvOptional, cuSeqlensCmpKvOptional, layoutKvOptional,	607	+ int64_t cmpKvBatchSize = GetCmpKvBatchSizeSmla(
549	- batchSize);	608	+ sequsedCmpKvOptional, cuSeqlensCmpKvOptional, layoutKvOptional, batchSize);
550	if (queryBatchSize != cmpKvBatchSize) {	609	if (queryBatchSize != cmpKvBatchSize) {
551	OP_LOGE(ACLNN_ERR_PARAM_INVALID, "when has_cmp_kv is true, the batch_size obtained from q should be "	610	OP_LOGE(ACLNN_ERR_PARAM_INVALID, "when has_cmp_kv is true, the batch_size obtained from q should be "
552	"the same as that obtained from cmp_kv, but got %lld and %lld", queryBatchSize,	611	"the same as that obtained from cmp_kv, but got %lld and %lld", queryBatchSize,
	@@ -577,6 +636,7 @@ static aclnnStatus ParamsCheck(const aclTensor *cuSeqlensQOptional, const aclTen
577	uint32_t aicCoreNum, uint32_t aivCoreNum, const char *socVersion,	636	uint32_t aicCoreNum, uint32_t aivCoreNum, const char *socVersion,
578	const aclTensor *metaData)	637	const aclTensor *metaData)
579	{	638	{
		639	+ bool isA5 = IsA5Smla(socVersion);
580	if (CheckSingleParamSmla(batchSize, maxSeqlenQ, maxSeqlenOriKv, maxSeqlenCmpKv, numHeadsQ, numHeadsKv, headDim,	640	if (CheckSingleParamSmla(batchSize, maxSeqlenQ, maxSeqlenOriKv, maxSeqlenCmpKv, numHeadsQ, numHeadsKv, headDim,
581	oriTopk, cmpTopk, cmpRatio, oriMaskMode, cmpMaskMode, oriWinLeft, oriWinRight,	641	oriTopk, cmpTopk, cmpRatio, oriMaskMode, cmpMaskMode, oriWinLeft, oriWinRight,
582	layoutQOptional, layoutKvOptional, hasOriKv, hasCmpKv, aicCoreNum, aivCoreNum,	642	layoutQOptional, layoutKvOptional, hasOriKv, hasCmpKv, aicCoreNum, aivCoreNum,
	@@ -588,7 +648,7 @@ static aclnnStatus ParamsCheck(const aclTensor *cuSeqlensQOptional, const aclTen
588	CheckConsistencySmla(cuSeqlensQOptional, cuSeqlensOriKvOptional, cuSeqlensCmpKvOptional, sequsedQOptional,	648	CheckConsistencySmla(cuSeqlensQOptional, cuSeqlensOriKvOptional, cuSeqlensCmpKvOptional, sequsedQOptional,
589	sequsedOriKvOptional, sequsedCmpKvOptional, cmpResidualKvOptional, oriTopkLengthOptional,	649	sequsedOriKvOptional, sequsedCmpKvOptional, cmpResidualKvOptional, oriTopkLengthOptional,
590	cmpTopkLengthOptional, batchSize, layoutQOptional, layoutKvOptional, hasOriKv, hasCmpKv,	650	cmpTopkLengthOptional, batchSize, layoutQOptional, layoutKvOptional, hasOriKv, hasCmpKv,
591	- metaData) == ACLNN_SUCCESS) {	651	+ isA5, metaData) == ACLNN_SUCCESS) {
592	return ACLNN_SUCCESS;	652	return ACLNN_SUCCESS;
593	} else {	653	} else {
594	return ACLNN_ERR_PARAM_INVALID;	654	return ACLNN_ERR_PARAM_INVALID;

Mattention/sparse_flash_mla_metadata/op_kernel_aicpu/sparse_flash_mla_metadata_aicpu.cpp+6-2

	@@ -278,6 +278,11 @@ bool SparseFlashMlaMetadataCpuKernel::ParamsInit()
278	CalcOriMaskMode();	278	CalcOriMaskMode();
279	CalcCmpMaskMode();	279	CalcCmpMaskMode();
280	isS1G_ = (layoutQ_ == "BSND" \|\| layoutQ_ == "BSH" \|\| layoutQ_ == "TND");	280	isS1G_ = (layoutQ_ == "BSND" \|\| layoutQ_ == "BSH" \|\| layoutQ_ == "TND");
		281	+ if (numHeadsKv_ == 0) {
		282	+ KERNEL_LOG_ERROR("num_heads_kv should not be 0.");
		283	+ return false;
		284	+ }
		285	+ ValidSocVersion validSocVersion = ProcessSocVersion();
281	groupSize_ = numHeadsQ_ / numHeadsKv_;	286	groupSize_ = numHeadsQ_ / numHeadsKv_;
282	if (hasOriKv_ && oriTopK_ != 0) {	287	if (hasOriKv_ && oriTopK_ != 0) {
283	isSparseOriKv_ = true;	288	isSparseOriKv_ = true;
	@@ -285,9 +290,8 @@ bool SparseFlashMlaMetadataCpuKernel::ParamsInit()
285	if (hasCmpKv_ && cmpTopK_ != 0) {	290	if (hasCmpKv_ && cmpTopK_ != 0) {
286	isSparseCmpKv_ = true;	291	isSparseCmpKv_ = true;
287	}	292	}
288	- ValidSocVersion validSocVersion = ProcessSocVersion();
289	if (validSocVersion == ValidSocVersion::ASCEND910) {	293	if (validSocVersion == ValidSocVersion::ASCEND910) {
290	- mBaseSize_ = groupSize_;	294	+ mBaseSize_ = isSparseCmpKv_ ? groupSize_ : (256U / groupSize_) * groupSize_;
291	s2BaseSize_ = 512U;	295	s2BaseSize_ = 512U;
292	} else if (validSocVersion == ValidSocVersion::ASCEND950) {	296	} else if (validSocVersion == ValidSocVersion::ASCEND950) {
293	if (numHeadsQ_ == 128) {	297	if (numHeadsQ_ == 128) {

Mattention/sparse_flash_mla_metadata/op_kernel_aicpu/sparse_flash_mla_metadata_aicpu.h+1-1

	@@ -372,4 +372,4 @@ private:
372	};	372	};
373	} // namespace aicpu	373	} // namespace aicpu
374		374
375	-#endif	375	+#endif

	@@ -56,13 +56,13 @@ using namespace SMLAKernel;
56	GET_TILING_DATA_WITH_STRUCT(tilingdataClass, tiling_data_in, tiling); \	56	GET_TILING_DATA_WITH_STRUCT(tilingdataClass, tiling_data_in, tiling); \
57	const tilingdataClass *__restrict tiling_data = &tiling_data_in; \	57	const tilingdataClass *__restrict tiling_data = &tiling_data_in; \
58	op.Init(query, oriKV, cmpKV, cmpSparseIndices, oriBlockTable, cmpBlockTable, cuSeqlensQ, \	58	op.Init(query, oriKV, cmpKV, cmpSparseIndices, oriBlockTable, cmpBlockTable, cuSeqlensQ, \
59	- cuSeqlensOriKv, cuSeqlensCmpKv, seqUsedQ, seqUsedOriKV, \	59	+ cuSeqlensOriKv, cuSeqlensCmpKv, seqUsedQ, seqUsedOriKV, seqUsedCmpKV, cmpResidualKV, \
60	sinks, metadata, attentionOut, softmaxLse, user, tiling_data, tiling, &tPipe); \	60	sinks, metadata, attentionOut, softmaxLse, user, tiling_data, tiling, &tPipe); \
61	op.Process(); \	61	op.Process(); \
62	} while (0)	62	} while (0)
63	#endif	63	#endif
64		64
65	-template <int FLASH_DECODE, int LAYOUT_T, int KV_LAYOUT_T, int TEMPLATE_MODE, int SPLIT_G>	65	+template <int FLASH_DECODE, int LAYOUT_T, int KV_LAYOUT_T, int TEMPLATE_MODE, int SPLIT_G, int HEAD_RATIO_ONE>
66	__global__ __aicore__ void	66	__global__ __aicore__ void
67	sparse_flash_mla(__gm__ uint8_t query, __gm__ uint8_t oriKV, __gm__ uint8_t *cmpKV,	67	sparse_flash_mla(__gm__ uint8_t query, __gm__ uint8_t oriKV, __gm__ uint8_t *cmpKV,
68	__gm__ uint8_t oriSparseIndices, __gm__ uint8_t cmpSparseIndices, __gm__ uint8_t *oriBlockTable,	68	__gm__ uint8_t oriSparseIndices, __gm__ uint8_t cmpSparseIndices, __gm__ uint8_t *oriBlockTable,
	@@ -92,22 +92,24 @@ sparse_flash_mla(__gm__ uint8_t query, __gm__ uint8_t oriKV, __gm__ uint8_t *c
92	if constexpr (ORIG_DTYPE_Q == DT_FLOAT16 && ORIG_DTYPE_ORI_KV == DT_FLOAT16 && ORIG_DTYPE_ATTN_OUT == DT_FLOAT16) {	92	if constexpr (ORIG_DTYPE_Q == DT_FLOAT16 && ORIG_DTYPE_ORI_KV == DT_FLOAT16 && ORIG_DTYPE_ATTN_OUT == DT_FLOAT16) {
93	if constexpr (TEMPLATE_MODE == SCFA_TEMPLATE) {	93	if constexpr (TEMPLATE_MODE == SCFA_TEMPLATE) {
94	SMLA_OP_IMPL(SparseFlashMlaScfa, SparseFlashMlaTilingData, half, half, half, FLASH_DECODE,	94	SMLA_OP_IMPL(SparseFlashMlaScfa, SparseFlashMlaTilingData, half, half, half, FLASH_DECODE,
95	- static_cast<SMLA_LAYOUT>(LAYOUT_T), static_cast<SMLA_LAYOUT>(KV_LAYOUT_T), TEMPLATE_MODE);	95	+ static_cast<SMLA_LAYOUT>(LAYOUT_T), static_cast<SMLA_LAYOUT>(KV_LAYOUT_T), TEMPLATE_MODE,
		96	+ static_cast<bool>(HEAD_RATIO_ONE));
96	} else {	97	} else {
97	SMLA_OP_IMPL(SparseFlashMlaSwa, SparseFlashMlaTilingData, half, half, half, FLASH_DECODE,	98	SMLA_OP_IMPL(SparseFlashMlaSwa, SparseFlashMlaTilingData, half, half, half, FLASH_DECODE,
98	- static_cast<SMLA_LAYOUT>(LAYOUT_T), static_cast<SMLA_LAYOUT>(KV_LAYOUT_T), TEMPLATE_MODE);	99	+ static_cast<SMLA_LAYOUT>(LAYOUT_T), static_cast<SMLA_LAYOUT>(KV_LAYOUT_T), TEMPLATE_MODE,
		100	+ static_cast<bool>(HEAD_RATIO_ONE));
99	}	101	}
100	}	102	}
101	if constexpr (ORIG_DTYPE_Q == DT_BF16 && ORIG_DTYPE_ORI_KV == DT_BF16 && ORIG_DTYPE_ATTN_OUT == DT_BF16) {	103	if constexpr (ORIG_DTYPE_Q == DT_BF16 && ORIG_DTYPE_ORI_KV == DT_BF16 && ORIG_DTYPE_ATTN_OUT == DT_BF16) {
102	if constexpr (TEMPLATE_MODE == SCFA_TEMPLATE) {	104	if constexpr (TEMPLATE_MODE == SCFA_TEMPLATE) {
103	SMLA_OP_IMPL(SparseFlashMlaScfa, SparseFlashMlaTilingData, bfloat16_t, bfloat16_t, bfloat16_t,	105	SMLA_OP_IMPL(SparseFlashMlaScfa, SparseFlashMlaTilingData, bfloat16_t, bfloat16_t, bfloat16_t,
104	FLASH_DECODE, static_cast<SMLA_LAYOUT>(LAYOUT_T), static_cast<SMLA_LAYOUT>(KV_LAYOUT_T),	106	FLASH_DECODE, static_cast<SMLA_LAYOUT>(LAYOUT_T), static_cast<SMLA_LAYOUT>(KV_LAYOUT_T),
105	- TEMPLATE_MODE);	107	+ TEMPLATE_MODE, static_cast<bool>(HEAD_RATIO_ONE));
106	} else {	108	} else {
107	SMLA_OP_IMPL(SparseFlashMlaSwa, SparseFlashMlaTilingData, bfloat16_t, bfloat16_t, bfloat16_t,	109	SMLA_OP_IMPL(SparseFlashMlaSwa, SparseFlashMlaTilingData, bfloat16_t, bfloat16_t, bfloat16_t,
108	FLASH_DECODE, static_cast<SMLA_LAYOUT>(LAYOUT_T), static_cast<SMLA_LAYOUT>(KV_LAYOUT_T),	110	FLASH_DECODE, static_cast<SMLA_LAYOUT>(LAYOUT_T), static_cast<SMLA_LAYOUT>(KV_LAYOUT_T),
109	- TEMPLATE_MODE);	111	+ TEMPLATE_MODE, static_cast<bool>(HEAD_RATIO_ONE));
110	}	112	}
111	}	113	}
112	#endif	114	#endif
113	-}	115	+}

	@@ -90,10 +90,8 @@ TEST_F(SparseFlashMlaTiling, test_tiling_swa_only_ori_kv_fp16_tnd_pa_nd)
90	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},	90	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},
91	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},	91	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
92	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},	92	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},
93	- {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BNBD")},	93	+ {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BBND")},
94	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},	94	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},
95	- {"ori_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
96	- {"cmp_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
97	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},	95	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},
98	},	96	},
99	&compileInfo, "Ascend910B", 40, 196608);	97	&compileInfo, "Ascend910B", 40, 196608);
	@@ -141,10 +139,8 @@ TEST_F(SparseFlashMlaTiling, test_tiling_swa_only_ori_kv_bf16_tnd_pa_nd)
141	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},	139	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},
142	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},	140	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
143	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},	141	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},
144	- {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BNBD")},	142	+ {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BBND")},
145	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},	143	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},
146	- {"ori_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
147	- {"cmp_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
148	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},	144	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},
149	},	145	},
150	&compileInfo, "Ascend910B", 40, 196608);	146	&compileInfo, "Ascend910B", 40, 196608);
	@@ -192,10 +188,8 @@ TEST_F(SparseFlashMlaTiling, test_tiling_cfa_ori_and_cmp_kv_fp16_tnd_pa_nd)
192	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},	188	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},
193	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},	189	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
194	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},	190	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},
195	- {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BNBD")},	191	+ {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BBND")},
196	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},	192	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},
197	- {"ori_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
198	- {"cmp_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
199	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},	193	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},
200	},	194	},
201	&compileInfo, "Ascend910B", 40, 196608);	195	&compileInfo, "Ascend910B", 40, 196608);
	@@ -243,10 +237,8 @@ TEST_F(SparseFlashMlaTiling, test_tiling_scfa_with_sparse_indices_fp16_tnd_pa_nd
243	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},	237	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},
244	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},	238	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
245	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},	239	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},
246	- {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BNBD")},	240	+ {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BBND")},
247	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},	241	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},
248	- {"ori_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
249	- {"cmp_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
250	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},	242	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},
251	},	243	},
252	&compileInfo, "Ascend910B", 40, 196608);	244	&compileInfo, "Ascend910B", 40, 196608);
	@@ -294,10 +286,8 @@ TEST_F(SparseFlashMlaTiling, test_tiling_scfa_with_sparse_indices_bf16_tnd_pa_nd
294	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},	286	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},
295	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},	287	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
296	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},	288	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},
297	- {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BNBD")},	289	+ {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BBND")},
298	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},	290	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},
299	- {"ori_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
300	- {"cmp_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
301	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},	291	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},
302	},	292	},
303	&compileInfo, "Ascend910B", 40, 196608);	293	&compileInfo, "Ascend910B", 40, 196608);
	@@ -345,10 +335,8 @@ TEST_F(SparseFlashMlaTiling, test_tiling_n1_not_64_failed)
345	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},	335	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},
346	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},	336	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
347	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},	337	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},
348	- {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BNBD")},	338	+ {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BBND")},
349	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},	339	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},
350	- {"ori_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
351	- {"cmp_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
352	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},	340	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},
353	},	341	},
354	&compileInfo, "Ascend910B", 40, 196608);	342	&compileInfo, "Ascend910B", 40, 196608);
	@@ -396,10 +384,8 @@ TEST_F(SparseFlashMlaTiling, test_tiling_ori_kv_null_failed)
396	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},	384	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},
397	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},	385	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
398	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},	386	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},
399	- {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BNBD")},	387	+ {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BBND")},
400	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},	388	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},
401	- {"ori_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
402	- {"cmp_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
403	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},	389	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},
404	},	390	},
405	&compileInfo, "Ascend910B", 40, 196608);	391	&compileInfo, "Ascend910B", 40, 196608);
	@@ -447,10 +433,8 @@ TEST_F(SparseFlashMlaTiling, test_tiling_cmp_sparse_indices_without_cmp_kv_faile
447	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},	433	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},
448	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},	434	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
449	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},	435	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},
450	- {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BNBD")},	436	+ {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BBND")},
451	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},	437	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},
452	- {"ori_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
453	- {"cmp_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
454	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},	438	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},
455	},	439	},
456	&compileInfo, "Ascend910B", 40, 196608);	440	&compileInfo, "Ascend910B", 40, 196608);
	@@ -498,10 +482,8 @@ TEST_F(SparseFlashMlaTiling, test_tiling_unsupported_dtype_failed)
498	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},	482	{"ori_win_left", Ops::Transformer::AnyValue::CreateFrom<int64_t>(127)},
499	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},	483	{"ori_win_right", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
500	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},	484	{"layout_q", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},
501	- {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BNBD")},	485	+ {"layout_kv", Ops::Transformer::AnyValue::CreateFrom<std::string>("PA_BBND")},
502	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},	486	{"topk_value_mode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},
503	- {"ori_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
504	- {"cmp_kv_stride", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
505	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},	487	{"return_softmax_lse", Ops::Transformer::AnyValue::CreateFrom<bool>(false)},
506	},	488	},
507	&compileInfo, "Ascend910B", 40, 196608);	489	&compileInfo, "Ascend910B", 40, 196608);