IFA算子新增sparsemode类型9，支持树形mask输入-ops-transformer-AtomGit

Mattention/common/op_host/arch32/fia_tiling_nonquant_mla.cpp+3-2

	@@ -137,8 +137,9 @@ bool FiaTilingNonQuantMla::IsCapable()
137	}	137	}
138		138
139	// 支持的sparse_mode值	139	// 支持的sparse_mode值
140	- if ((fiaInfo_->sparseMode != SPARSE_MODE_NO_MASK) && (fiaInfo_->sparseMode != SPARSE_MODE_RIGHT_DOWN) &&	140	+
141	- (fiaInfo_->sparseMode != SPARSE_MODE_BAND)) {	141	+ if ((fiaInfo_->sparseMode != SPARSE_MODE_NO_MASK) && (fiaInfo_->sparseMode != SPARSE_MODE_RIGHT_DOWN) &&
		142	+ (fiaInfo_->sparseMode != SPARSE_MODE_BAND) && (fiaInfo_->sparseMode != SPARSE_MODE_TREE)) {
142	return false;	143	return false;
143	}	144	}
144		145

Mattention/common/op_host/fia_tiling_info.h+2-0文件内容审核中，请稍后刷新重试

Mattention/common/op_host/fia_tiling_shape.cpp+3-0

	@@ -39,6 +39,7 @@ static const std::map<FiaLayout, std::vector<FiaAxis>> FIA_LAYOUT_AXIS_MAP = {
39	{FiaLayout::B1S1S2, {FiaAxis::B, FiaAxis::CONST, FiaAxis::S1, FiaAxis::S2}},	39	{FiaLayout::B1S1S2, {FiaAxis::B, FiaAxis::CONST, FiaAxis::S1, FiaAxis::S2}},
40	{FiaLayout::IS1S2, {FiaAxis::CONST, FiaAxis::S1, FiaAxis::S2}},	40	{FiaLayout::IS1S2, {FiaAxis::CONST, FiaAxis::S1, FiaAxis::S2}},
41	{FiaLayout::I1S1S2, {FiaAxis::CONST, FiaAxis::CONST, FiaAxis::S1, FiaAxis::S2}},	41	{FiaLayout::I1S1S2, {FiaAxis::CONST, FiaAxis::CONST, FiaAxis::S1, FiaAxis::S2}},
		42	+ {FiaLayout::S1S1, {FiaAxis::S1}},
42	};	43	};
43		44
44	static bool equal_to(const int64_t& a, const int64_t& b)	45	static bool equal_to(const int64_t& a, const int64_t& b)
	@@ -243,6 +244,8 @@ ge::graphStatus FiaTilingShapeCompare::GetExpectedShapeSpecial(gert::Shape &shap
243	shapeExpected = gert::Shape({param.CONST, param.S1, param.S2});	244	shapeExpected = gert::Shape({param.CONST, param.S1, param.S2});
244	} else if (layout_ == FiaLayout::I1S1S2) {	245	} else if (layout_ == FiaLayout::I1S1S2) {
245	shapeExpected = gert::Shape({param.CONST, param.CONST, param.S1, param.S2});	246	shapeExpected = gert::Shape({param.CONST, param.CONST, param.S1, param.S2});
		247	+ } else if (layout_ == FiaLayout::S1S1) {
		248	+ shapeExpected = gert::Shape({param.S1});
246	} else {	249	} else {
247	OP_LOGE(opName_, "[%s] layout %s is unsupported", funcName.c_str(), LayoutToSerialString(layout_).c_str());	250	OP_LOGE(opName_, "[%s] layout %s is unsupported", funcName.c_str(), LayoutToSerialString(layout_).c_str());
248	return ge::GRAPH_FAILED;	251	return ge::GRAPH_FAILED;

Mattention/common/op_host/split_core.cpp+2-0文件内容审核中，请稍后刷新重试

Mattention/common/op_host/split_core.h+1-0文件内容审核中，请稍后刷新重试

Mattention/common/op_kernel/arch32/fia_block_vec_nonquant.h+33-10

	@@ -579,7 +579,6 @@ __aicore__ inline void FiaBlockVecNonQuant<FIAT>::ElewiseCompute(
579	maskInfo.maskValue = negativeIntScalar;	579	maskInfo.maskValue = negativeIntScalar;
580	maskInfo.s1LeftPaddingSize = info.qPaddingBeginOffset;	580	maskInfo.s1LeftPaddingSize = info.qPaddingBeginOffset;
581	maskInfo.s2LeftPaddingSize = info.kvPaddingBeginOffset;	581	maskInfo.s2LeftPaddingSize = info.kvPaddingBeginOffset;
582	-
583	if (constInfo.qSeqSize == 1) {	582	if (constInfo.qSeqSize == 1) {
584	maskInfo.layout = fa_base_vector::S1_EQUAL1;	583	maskInfo.layout = fa_base_vector::S1_EQUAL1;
585	} else if constexpr (LAYOUT_T == FIA_LAYOUT::TND \|\| LAYOUT_T == FIA_LAYOUT::BSH) {	584	} else if constexpr (LAYOUT_T == FIA_LAYOUT::TND \|\| LAYOUT_T == FIA_LAYOUT::BSH) {
	@@ -587,25 +586,49 @@ __aicore__ inline void FiaBlockVecNonQuant<FIAT>::ElewiseCompute(
587	} else {	586	} else {
588	maskInfo.layout = fa_base_vector::GS;	587	maskInfo.layout = fa_base_vector::GS;
589	}	588	}
		589	+
590	maskInfo.attenMaskType = fa_base_vector::MASK_BOOL; // compatible with int8/uint8	590	maskInfo.attenMaskType = fa_base_vector::MASK_BOOL; // compatible with int8/uint8
591	- LocalTensor<bool> maskUb;	591	+
592	- LocalTensor<bool> attenMaskTmpUb;	592	+ // 添加Sparse9的处理，由于sparse9的mask拷贝只占最小块的一部分，所以需要对UB空间赋初值0，表示不被掩码覆盖
		593	+ // TND场景下mask传入∑s1²，其余场景传入[B,S1,S1]
		594	+ LocalTensor<bool> maskUb = inputQue2.AllocTensor<bool>();
		595	+ LocalTensor<bool> attenMaskTmpUb = maskUb[BUFFER_SIZE_BYTE_16K / 2];
593	LocalTensor<uint8_t> ubWorkSpace = tmpBuf.Get<uint8_t>();	596	LocalTensor<uint8_t> ubWorkSpace = tmpBuf.Get<uint8_t>();
		597	+ event_t eventIdVMte2 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2));
		598	+ if (maskInfo.sparseMode == fa_base_vector::TREE) {
		599	+ LocalTensor<int16_t> mask16 = maskUb.template ReinterpretCast<int16_t>();
		600	+ uint32_t zeroCount = BUFFER_SIZE_BYTE_8K / sizeof(int16_t);
		601	+ Duplicate(mask16, static_cast<int16_t>(0), zeroCount);
		602	+ maskUb = mask16.template ReinterpretCast<bool>();
		603	+ SetFlag<HardEvent::V_MTE2>(eventIdVMte2);
		604	+ WaitFlag<HardEvent::V_MTE2>(eventIdVMte2);
		605	+ // 修改attenMaskStride、attenMaskBatchStride值
		606	+ maskInfo.attenMaskBatchStride = maskInfo.attenMaskBatchStride * maskInfo.batchIdx;
		607	+ if (LAYOUT_T == FIA_LAYOUT::TND \|\| LAYOUT_T == FIA_LAYOUT::NTD) {
		608	+ maskInfo.attenMaskStride = info.actS1Size;
		609	+ maskInfo.attenMaskBatchStride = 0;
		610	+ for (int32_t i = 0; i < maskInfo.batchIdx; i++) {
		611	+ maskInfo.attenMaskBatchStride += qActSeqLensParser.GetActualSeqLength(i) * qActSeqLensParser.GetActualSeqLength(i);
		612	+ }
		613	+ }
		614	+ }
		615	+
594	if (!fa_base_vector::IsSkipAttentionmask(maskInfo)) {	616	if (!fa_base_vector::IsSkipAttentionmask(maskInfo)) {
595	- maskUb = inputQue2.AllocTensor<bool>();	617	+ if (maskInfo.sparseMode == fa_base_vector::TREE) {
596	- attenMaskTmpUb = maskUb[BUFFER_SIZE_BYTE_16K / 2];	618	+ fa_base_vector::AttentionmaskCopyIn<bool, bool, true>(maskUb, attenMaskBoolGm, attenMaskTmpUb, maskInfo);
597	- fa_base_vector::AttentionmaskCopyIn(maskUb, attenMaskBoolGm, attenMaskTmpUb, maskInfo);	619	+ } else {
		620	+ fa_base_vector::AttentionmaskCopyIn(maskUb, attenMaskBoolGm, attenMaskTmpUb, maskInfo);
		621	+ }
598	AscendC::PipeBarrier<PIPE_V>();	622	AscendC::PipeBarrier<PIPE_V>();
599	fa_base_vector::AttentionMaskCompute<MM1_OUT_T>(mmResUb, mmResUb, maskUb, ubWorkSpace, maskInfo);	623	fa_base_vector::AttentionMaskCompute<MM1_OUT_T>(mmResUb, mmResUb, maskUb, ubWorkSpace, maskInfo);
600	- inputQue2.FreeTensor(maskUb);
601	}	624	}
602	if (!fa_base_vector::IsSkipAttentionmaskForPre(maskInfo)) {	625	if (!fa_base_vector::IsSkipAttentionmaskForPre(maskInfo)) {
603	- maskUb = inputQue2.AllocTensor<bool>();	626	+ SetFlag<HardEvent::V_MTE2>(eventIdVMte2);
604	- attenMaskTmpUb = maskUb[BUFFER_SIZE_BYTE_16K / 2];	627	+ WaitFlag<HardEvent::V_MTE2>(eventIdVMte2);
605	fa_base_vector::AttentionmaskCopyIn(maskUb, attenMaskBoolGm, attenMaskTmpUb, maskInfo, true);	628	fa_base_vector::AttentionmaskCopyIn(maskUb, attenMaskBoolGm, attenMaskTmpUb, maskInfo, true);
606	fa_base_vector::AttentionMaskCompute<MM1_OUT_T>(mmResUb, mmResUb, maskUb, ubWorkSpace, maskInfo, true);	629	fa_base_vector::AttentionMaskCompute<MM1_OUT_T>(mmResUb, mmResUb, maskUb, ubWorkSpace, maskInfo, true);
607	- inputQue2.FreeTensor(maskUb);
608	}	630	}
		631	+ inputQue2.FreeTensor(maskUb);
609	}	632	}
610	}	633	}
611		634

Mattention/common/op_kernel/arch32/fia_block_vec_nonquant_mla.h+24-1

	@@ -454,13 +454,36 @@ __aicore__ inline void FiaBlockVecNonQuantMla<FIAT>::ElewiseCompute(
454	}	454	}
455		455
456	maskInfo.attenMaskType = fa_base_vector::MASK_BOOL; // compatible with int8/uint8	456	maskInfo.attenMaskType = fa_base_vector::MASK_BOOL; // compatible with int8/uint8
		457	+
		458	+ // 添加Sparse9的处理，由于sparse9的mask拷贝只占最小块的一部分，所以需要对UB空间赋初值0，表示不被掩码覆盖
		459	+ // TND场景下mask传入∑s1²，其余场景传入[B,S1,S1]
457	LocalTensor<bool> maskUb = inputBuff2.Get<bool>();	460	LocalTensor<bool> maskUb = inputBuff2.Get<bool>();
458	maskUb = maskUb[pingpongFlag * INPUT2_BUFFER_OFFSET / sizeof(bool)];	461	maskUb = maskUb[pingpongFlag * INPUT2_BUFFER_OFFSET / sizeof(bool)];
		462	+ if (maskInfo.sparseMode == fa_base_vector::TREE) {
		463	+ LocalTensor<int16_t> mask16 = maskUb.template ReinterpretCast<int16_t>();
		464	+ AscendC::Duplicate(mask16, static_cast<int16_t>(0), INPUT2_BUFFER_OFFSET / sizeof(int16_t));
		465	+ maskUb = mask16.template ReinterpretCast<bool>();
		466	+ WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF2_FLAG + pingpongFlag);
		467	+ SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF2_FLAG + pingpongFlag);
		468	+ // 修改attenMaskStride、attenMaskBatchStride值
		469	+ maskInfo.attenMaskBatchStride = maskInfo.attenMaskBatchStride * maskInfo.batchIdx;
		470	+ if (LAYOUT_T == FIA_LAYOUT::TND \|\| LAYOUT_T == FIA_LAYOUT::NTD) {
		471	+ maskInfo.attenMaskStride = info.actS1Size;
		472	+ maskInfo.attenMaskBatchStride = 0;
		473	+ for (int32_t i = 0; i < maskInfo.batchIdx; i++) {
		474	+ maskInfo.attenMaskBatchStride += qActSeqLensParser.GetActualSeqLength(i) * qActSeqLensParser.GetActualSeqLength(i);
		475	+ }
		476	+ }
		477	+ }
459	LocalTensor<bool> attenMaskTmpUb = attenMaskTmpBuff.Get<bool>();	478	LocalTensor<bool> attenMaskTmpUb = attenMaskTmpBuff.Get<bool>();
460	LocalTensor<uint8_t> ubWorkSpace = tmpBuf.Get<uint8_t>();	479	LocalTensor<uint8_t> ubWorkSpace = tmpBuf.Get<uint8_t>();
461	if (!fa_base_vector::IsSkipAttentionmask(maskInfo)) {	480	if (!fa_base_vector::IsSkipAttentionmask(maskInfo)) {
462	WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF2_FLAG + pingpongFlag);	481	WaitFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF2_FLAG + pingpongFlag);
463	- fa_base_vector::AttentionmaskCopyIn(maskUb, attenMaskBoolGm, attenMaskTmpUb, maskInfo);	482	+ if (maskInfo.sparseMode == fa_base_vector::TREE) {
		483	+ fa_base_vector::AttentionmaskCopyIn<bool, bool, true>(maskUb, attenMaskBoolGm, attenMaskTmpUb, maskInfo);
		484	+ } else {
		485	+ fa_base_vector::AttentionmaskCopyIn(maskUb, attenMaskBoolGm, attenMaskTmpUb, maskInfo);
		486	+ }
464	AscendC::PipeBarrier<PIPE_V>();	487	AscendC::PipeBarrier<PIPE_V>();
465	fa_base_vector::AttentionMaskCompute<MM1_OUT_T>(mmResUb, mmResUb, maskUb, ubWorkSpace, maskInfo);	488	fa_base_vector::AttentionMaskCompute<MM1_OUT_T>(mmResUb, mmResUb, maskUb, ubWorkSpace, maskInfo);
466	SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF2_FLAG + pingpongFlag);	489	SetFlag<AscendC::HardEvent::V_MTE2>(SYNC_INPUT_BUF2_FLAG + pingpongFlag);

Mattention/common/op_kernel/arch32/fia_kernel_nonquant.h+1-1文件内容审核中，请稍后刷新重试

Mattention/common/op_kernel/arch32/fia_kernel_nonquant_mla.h+1-1文件内容审核中，请稍后刷新重试

Mattention/common/op_kernel/vector_common.h+121-26

	@@ -47,6 +47,7 @@ enum SparseMode : uint8_t {
47	LEFT_UP_CAUSAL,	47	LEFT_UP_CAUSAL,
48	RIGHT_DOWN_CAUSAL,	48	RIGHT_DOWN_CAUSAL,
49	BAND,	49	BAND,
		50	+ TREE = 9,
50	};	51	};
51		52
52	__aicore__ inline bool IsExistInvalidRows(int64_t nextTokensPerBatch, int64_t preTokensPerBatch, uint32_t mode,	53	__aicore__ inline bool IsExistInvalidRows(int64_t nextTokensPerBatch, int64_t preTokensPerBatch, uint32_t mode,
	@@ -853,6 +854,14 @@ __aicore__ inline uint64_t ComputeAttenMaskOffsetNoCompress(MaskInfo &info, uint
853	return bOffset + s1Offset + s2Offset;	854	return bOffset + s1Offset + s2Offset;
854	}	855	}
855		856
		857	+__aicore__ inline uint64_t ComputeAttenMaskOffsetTree(MaskInfo &info, uint32_t s1StartIdx, uint64_t treeMaskStart)
		858	+{
		859	+ uint64_t bOffset = info.attenMaskBatchStride;
		860	+ uint64_t s1Offset = (s1StartIdx % info.s1Size) * info.attenMaskStride;
		861	+ uint64_t s2Offset = info.s2StartIdx > treeMaskStart ? info.s2StartIdx - treeMaskStart : 0;
		862	+ return bOffset + s1Offset + s2Offset;
		863	+}
		864	+
856	__aicore__ inline uint64_t ComputeAttenMaskOffsetCompress(MaskInfo &info, uint32_t s1StartIdx)	865	__aicore__ inline uint64_t ComputeAttenMaskOffsetCompress(MaskInfo &info, uint32_t s1StartIdx)
857	{	866	{
858	int64_t nextToken = 0; // sparse2 本身原点就是左上角	867	int64_t nextToken = 0; // sparse2 本身原点就是左上角
	@@ -886,24 +895,58 @@ __aicore__ inline uint64_t ComputeAttenMaskOffsetCompressPre(MaskInfo &info, uin
886	return offset;	895	return offset;
887	}	896	}
888		897
889	-__aicore__ inline uint64_t ComputeAttenMaskOffset(MaskInfo &info, uint32_t s1StartIdx = 0, bool isPre = false)	898	+template <bool ENABLE_TREE = false>
		899	+__aicore__ inline uint64_t ComputeAttenMaskOffset(MaskInfo &info, uint32_t s1StartIdx = 0, uint64_t treeMaskStart = 0, bool isPre = false)
890	{	900	{
891	if (isPre) {	901	if (isPre) {
892	return ComputeAttenMaskOffsetCompressPre(info, s1StartIdx);	902	return ComputeAttenMaskOffsetCompressPre(info, s1StartIdx);
893	} else {	903	} else {
894	if (info.sparseMode == DEFAULT_MASK \|\| info.sparseMode == ALL_MASK) {	904	if (info.sparseMode == DEFAULT_MASK \|\| info.sparseMode == ALL_MASK) {
895	return ComputeAttenMaskOffsetNoCompress(info, s1StartIdx);	905	return ComputeAttenMaskOffsetNoCompress(info, s1StartIdx);
896	- } else {
897	- return ComputeAttenMaskOffsetCompress(info, s1StartIdx);
898	}	906	}
		907	+ if constexpr (ENABLE_TREE) {
		908	+ if (info.sparseMode == TREE) {
		909	+ return ComputeAttenMaskOffsetTree(info, s1StartIdx, treeMaskStart);
		910	+ }
		911	+ }
		912	+ return ComputeAttenMaskOffsetCompress(info, s1StartIdx);
899	}	913	}
900	}	914	}
901		915
902	-template <typename T>	916	+template <typename T, bool ENABLE_TREE = false>
903	__aicore__ inline void AttentionmaskDataCopy(LocalTensor<T> &attenMaskUb, GlobalTensor<T> &srcGmAddr, MaskInfo &info, uint32_t s1StartIdx, uint32_t s1EndIdx, bool isPre = false)	917	__aicore__ inline void AttentionmaskDataCopy(LocalTensor<T> &attenMaskUb, GlobalTensor<T> &srcGmAddr, MaskInfo &info, uint32_t s1StartIdx, uint32_t s1EndIdx, bool isPre = false)
904	{	918	{
		919	+ uint32_t treeMaskStart = 0;
		920	+ if constexpr (ENABLE_TREE) {
		921	+ treeMaskStart = info.s2Size - info.s1Size;
		922	+ uint32_t curS2EnsPos = info.s2StartIdx + info.s2dealNum;
		923	+ // 只有info.s2StartIdx + info.s2dealNum > treeMaskStart时，才会进入此流程；
		924	+ // 当info.s2StartIdx > treeMaskStart，mask拷贝也是全量拷贝，和其余sparse过程相同
		925	+ if (info.s2StartIdx < treeMaskStart) {
		926	+ if (curS2EnsPos <= treeMaskStart) {
		927	+ return; // 整个 tile 在零区，UB 已初始化为 0，无需拷贝
		928	+ }
		929	+ // 由于sparse9 mask只有一部分，不能合并处理
		930	+ uint32_t attenMaskSize = curS2EnsPos - treeMaskStart;
		931	+ uint32_t attenMaskSizeAlign = Align(static_cast<uint32_t>(attenMaskSize + treeMaskStart % 32), 32U);
		932	+ uint64_t maskOffset = ComputeAttenMaskOffset<ENABLE_TREE>(info, s1StartIdx, treeMaskStart, isPre);
		933	+ DataCopyExtParams dataCopyParams;
		934	+ dataCopyParams.blockCount = s1EndIdx - s1StartIdx;
		935	+ dataCopyParams.blockLen = curS2EnsPos - treeMaskStart;
		936	+ dataCopyParams.srcStride = info.attenMaskStride - (curS2EnsPos - treeMaskStart);
		937	+ dataCopyParams.dstStride = (treeMaskStart - info.s2StartIdx) / 32; // dst在UB上，单位为32字节，同时左侧不对齐场景会进行左padding
		938	+ DataCopyPadExtParams<bool> padParams;
		939	+ padParams.isPad = true;
		940	+ padParams.leftPadding = static_cast<uint8_t>(treeMaskStart % 32);
		941	+ padParams.rightPadding = static_cast<uint8_t>(attenMaskSizeAlign - (attenMaskSize + treeMaskStart % 32));
		942	+ padParams.paddingValue = 0;
		943	+ DataCopyPad(attenMaskUb[(treeMaskStart - info.s2StartIdx) / 32 * 32], srcGmAddr[maskOffset], dataCopyParams, padParams);
		944	+ return;
		945	+ }
		946	+ }
		947	+ // 标准拷贝路径（非TREE模式，或TREE模式但s2StartIdx >= treeMaskStart时）
905	uint32_t attenMaskSizeAlign = Align(info.s2dealNum, 32U);	948	uint32_t attenMaskSizeAlign = Align(info.s2dealNum, 32U);
906	- uint64_t maskOffset = ComputeAttenMaskOffset(info, s1StartIdx, isPre);	949	+ uint64_t maskOffset = ComputeAttenMaskOffset<ENABLE_TREE>(info, s1StartIdx, treeMaskStart, isPre);
907	DataCopyExtParams dataCopyParams;	950	DataCopyExtParams dataCopyParams;
908	dataCopyParams.blockCount = s1EndIdx - s1StartIdx;	951	dataCopyParams.blockCount = s1EndIdx - s1StartIdx;
909	dataCopyParams.blockLen = info.s2dealNum;	952	dataCopyParams.blockLen = info.s2dealNum;
	@@ -914,7 +957,7 @@ __aicore__ inline void AttentionmaskDataCopy(LocalTensor<T> &attenMaskUb, Global
914	DataCopyPad(attenMaskUb, srcGmAddr[maskOffset], dataCopyParams, padParams);	957	DataCopyPad(attenMaskUb, srcGmAddr[maskOffset], dataCopyParams, padParams);
915	}	958	}
916		959
917	-template <typename T, typename U>	960	+template <typename T, typename U, bool ENABLE_TREE = false>
918	__aicore__ inline void AttentionmaskCopyInForGsLayout(LocalTensor<T> &attenMaskUb, GlobalTensor<T> &srcGmAddr, LocalTensor<U> &tmpBuf, MaskInfo &info, bool isPre = false)	961	__aicore__ inline void AttentionmaskCopyInForGsLayout(LocalTensor<T> &attenMaskUb, GlobalTensor<T> &srcGmAddr, LocalTensor<U> &tmpBuf, MaskInfo &info, bool isPre = false)
919	{	962	{
920	int32_t s1StartIdx = info.gs1StartIdx % info.s1Size;	963	int32_t s1StartIdx = info.gs1StartIdx % info.s1Size;
	@@ -922,17 +965,17 @@ __aicore__ inline void AttentionmaskCopyInForGsLayout(LocalTensor<T> &attenMaskU
922	uint32_t attenMaskSizeAlign = Align(info.s2dealNum, 32U);	965	uint32_t attenMaskSizeAlign = Align(info.s2dealNum, 32U);
923	if (info.gs1dealNum <= info.s1Size) {	966	if (info.gs1dealNum <= info.s1Size) {
924	if (s1StartIdx + info.gs1dealNum > info.s1Size) {	967	if (s1StartIdx + info.gs1dealNum > info.s1Size) {
925	- AttentionmaskDataCopy(attenMaskUb, srcGmAddr, info, s1StartIdx, info.s1Size, isPre);	968	+ AttentionmaskDataCopy<T, ENABLE_TREE>(attenMaskUb, srcGmAddr, info, s1StartIdx, info.s1Size, isPre);
926	LocalTensor<T> attenMaskSecUb = attenMaskUb[(info.s1Size - s1StartIdx) * attenMaskSizeAlign];	969	LocalTensor<T> attenMaskSecUb = attenMaskUb[(info.s1Size - s1StartIdx) * attenMaskSizeAlign];
927	- AttentionmaskDataCopy(attenMaskSecUb, srcGmAddr, info, 0, s1EndIdx, isPre);	970	+ AttentionmaskDataCopy<T, ENABLE_TREE>(attenMaskSecUb, srcGmAddr, info, 0, s1EndIdx, isPre);
928	} else {	971	} else {
929	- AttentionmaskDataCopy(attenMaskUb, srcGmAddr, info, s1StartIdx, s1EndIdx, isPre);	972	+ AttentionmaskDataCopy<T, ENABLE_TREE>(attenMaskUb, srcGmAddr, info, s1StartIdx, s1EndIdx, isPre);
930	}	973	}
931	event_t enQueEvtID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));	974	event_t enQueEvtID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
932	SetFlag<HardEvent::MTE2_V>(enQueEvtID);	975	SetFlag<HardEvent::MTE2_V>(enQueEvtID);
933	WaitFlag<HardEvent::MTE2_V>(enQueEvtID);	976	WaitFlag<HardEvent::MTE2_V>(enQueEvtID);
934	} else {	977	} else {
935	- AttentionmaskDataCopy(attenMaskUb, srcGmAddr, info, 0, info.s1Size, isPre);	978	+ AttentionmaskDataCopy<T, ENABLE_TREE>(attenMaskUb, srcGmAddr, info, 0, info.s1Size, isPre);
936	LocalTensor<T> attenMaskUbDst = tmpBuf.template ReinterpretCast<T>();	979	LocalTensor<T> attenMaskUbDst = tmpBuf.template ReinterpretCast<T>();
937	event_t enQueEvtID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));	980	event_t enQueEvtID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
938	SetFlag<HardEvent::MTE2_V>(enQueEvtID);	981	SetFlag<HardEvent::MTE2_V>(enQueEvtID);
	@@ -964,7 +1007,7 @@ __aicore__ inline void AttentionmaskCopyInForGsLayout(LocalTensor<T> &attenMaskU
964	}	1007	}
965	}	1008	}
966		1009
967	-template <typename T, typename U>	1010	+template <typename T, typename U, bool ENABLE_TREE = false>
968	__aicore__ inline void AttentionmaskCopyInForSgLayout(LocalTensor<T> &attenMaskUb, GlobalTensor<T> &srcGmAddr, LocalTensor<U> &tmpBuf, MaskInfo &info, bool isPre = false)	1011	__aicore__ inline void AttentionmaskCopyInForSgLayout(LocalTensor<T> &attenMaskUb, GlobalTensor<T> &srcGmAddr, LocalTensor<U> &tmpBuf, MaskInfo &info, bool isPre = false)
969	{	1012	{
970	uint32_t s1StartIdx = info.gs1StartIdx / info.gSize;	1013	uint32_t s1StartIdx = info.gs1StartIdx / info.gSize;
	@@ -972,7 +1015,7 @@ __aicore__ inline void AttentionmaskCopyInForSgLayout(LocalTensor<T> &attenMaskU
972	uint32_t s1Count = s1EndIdx - s1StartIdx + 1;	1015	uint32_t s1Count = s1EndIdx - s1StartIdx + 1;
973	uint32_t attenMaskSizeAlign = Align(info.s2dealNum, 32U);	1016	uint32_t attenMaskSizeAlign = Align(info.s2dealNum, 32U);
974		1017
975	- AttentionmaskDataCopy(attenMaskUb, srcGmAddr, info, s1StartIdx, s1EndIdx + 1, isPre);	1018	+ AttentionmaskDataCopy<T, ENABLE_TREE>(attenMaskUb, srcGmAddr, info, s1StartIdx, s1EndIdx + 1, isPre);
976		1019
977	event_t enQueEvtID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));	1020	event_t enQueEvtID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
978	SetFlag<HardEvent::MTE2_V>(enQueEvtID);	1021	SetFlag<HardEvent::MTE2_V>(enQueEvtID);
	@@ -1012,12 +1055,25 @@ __aicore__ inline void AttentionmaskCopyInForSgLayout(LocalTensor<T> &attenMaskU
1012	attenMaskUb = attenMaskUbDst.template ReinterpretCast<bool>();	1055	attenMaskUb = attenMaskUbDst.template ReinterpretCast<bool>();
1013	}	1056	}
1014		1057
		1058	+template <bool ENABLE_TREE = false>
1015	__aicore__ inline bool IsSkipAttentionmask(MaskInfo &info)	1059	__aicore__ inline bool IsSkipAttentionmask(MaskInfo &info)
1016	{	1060	{
1017	if (info.sparseMode == DEFAULT_MASK \|\| info.sparseMode == ALL_MASK) {	1061	if (info.sparseMode == DEFAULT_MASK \|\| info.sparseMode == ALL_MASK) {
1018	return false;	1062	return false;
1019	}	1063	}
1020		1064
		1065	+ // 增加sparse = 9的处理
		1066	+ if constexpr (ENABLE_TREE) {
		1067	+ if (info.sparseMode == TREE) {
		1068	+ // 由于分核时按照Batch进行划分，sparse9在每个batch的所有 S 跳过的范围固定，所以不区分跨g轴的情况
		1069	+ if (static_cast<int64_t>(info.s2StartIdx + info.s2dealNum) > static_cast<int64_t>(info.s2Size - info.s1Size)) {
		1070	+ return false;
		1071	+ } else {
		1072	+ return true;
		1073	+ }
		1074	+ }
		1075	+ }
		1076	+
1021	int32_t s1StartIdx = info.layout == GS ? info.gs1StartIdx % info.s1Size : info.gs1StartIdx / info.gSize;	1077	int32_t s1StartIdx = info.layout == GS ? info.gs1StartIdx % info.s1Size : info.gs1StartIdx / info.gSize;
1022	if (info.layout == GS && s1StartIdx + info.gs1dealNum > info.s1Size) { // 当跨多个s1时，不再支持跳过计算	1078	if (info.layout == GS && s1StartIdx + info.gs1dealNum > info.s1Size) { // 当跨多个s1时，不再支持跳过计算
1023	return false;	1079	return false;
	@@ -1056,30 +1112,69 @@ __aicore__ inline bool IsSkipAttentionmaskForPre(MaskInfo &info)
1056	return false;	1112	return false;
1057	}	1113	}
1058		1114
1059	-template <typename T, typename U>	1115	+template <typename T, typename U, bool ENABLE_TREE = false>
1060	__aicore__ inline void AttentionmaskCopyIn(LocalTensor<T> &attenMaskUb, GlobalTensor<T> &srcGmAddr, LocalTensor<U> &tmpBuf, MaskInfo &info, bool isPre = false)	1116	__aicore__ inline void AttentionmaskCopyIn(LocalTensor<T> &attenMaskUb, GlobalTensor<T> &srcGmAddr, LocalTensor<U> &tmpBuf, MaskInfo &info, bool isPre = false)
1061	{	1117	{
1062	if (info.layout == GS) {	1118	if (info.layout == GS) {
1063	- AttentionmaskCopyInForGsLayout(attenMaskUb, srcGmAddr, tmpBuf, info, isPre);	1119	+ AttentionmaskCopyInForGsLayout<T, U, ENABLE_TREE>(attenMaskUb, srcGmAddr, tmpBuf, info, isPre);
1064	} else if(info.layout == SG) { // sg	1120	} else if(info.layout == SG) { // sg
1065	- AttentionmaskCopyInForSgLayout(attenMaskUb, srcGmAddr, tmpBuf, info, isPre);	1121	+ AttentionmaskCopyInForSgLayout<T, U, ENABLE_TREE>(attenMaskUb, srcGmAddr, tmpBuf, info, isPre);
1066	} else if (info.layout == S1_EQUAL1) {	1122	} else if (info.layout == S1_EQUAL1) {
1067	- uint64_t maskOffset = ComputeAttenMaskOffset(info, 0, isPre);	1123	+ uint32_t treeMaskStart = 0;
1068	- uint32_t attenMaskSizeAlign = Align(info.s2dealNum, 32U);	1124	+ if constexpr (ENABLE_TREE) {
1069	- DataCopyExtParams dataCopyParams;	1125	+ treeMaskStart = info.s2Size - info.s1Size;
1070	- dataCopyParams.blockCount = 1;	1126	+ uint32_t curS2EndPos = info.s2StartIdx + info.s2dealNum;
1071	- dataCopyParams.blockLen = info.s2dealNum;	1127	+ if (info.s2StartIdx >= treeMaskStart) {
1072	- dataCopyParams.srcStride = info.attenMaskStride - info.s2dealNum ;	1128	+ // TREE 但 s2StartIdx >= treeMaskStart：全量拷贝
1073	- dataCopyParams.dstStride = 0;	1129	+ uint64_t maskOffset = ComputeAttenMaskOffset<ENABLE_TREE>(info, 0, treeMaskStart, isPre);
1074	- DataCopyPadExtParams<bool> padParams{true, 0, static_cast<uint8_t>(attenMaskSizeAlign - info.s2dealNum), 0};	1130	+ uint32_t attenMaskSizeAlign = Align(info.s2dealNum, 32U);
1075	- DataCopyPad(attenMaskUb, srcGmAddr[maskOffset], dataCopyParams, padParams);	1131	+ DataCopyExtParams dataCopyParams;
		1132	+ dataCopyParams.blockCount = 1;
		1133	+ dataCopyParams.blockLen = info.s2dealNum;
		1134	+ dataCopyParams.srcStride = info.attenMaskStride - info.s2dealNum;
		1135	+ dataCopyParams.dstStride = 0;
		1136	+ DataCopyPadExtParams<bool> padParams{true, 0,
		1137	+ static_cast<uint8_t>(attenMaskSizeAlign - info.s2dealNum), 0};
		1138	+ DataCopyPad(attenMaskUb, srcGmAddr[maskOffset], dataCopyParams, padParams);
		1139	+ } else if (curS2EndPos > treeMaskStart) {
		1140	+ // 部分拷贝：只拷贝 [treeMaskStart, curS2EndPos) 区域
		1141	+ uint32_t attenMaskSize = curS2EndPos - treeMaskStart;
		1142	+ uint32_t attenMaskSizeAlign = Align(static_cast<uint32_t>(attenMaskSize + treeMaskStart % 32), 32U);
		1143	+ uint64_t maskOffset = ComputeAttenMaskOffset<ENABLE_TREE>(info, 0, treeMaskStart, isPre);
		1144	+ DataCopyExtParams dataCopyParams;
		1145	+ dataCopyParams.blockCount = 1;
		1146	+ dataCopyParams.blockLen = curS2EndPos - treeMaskStart;
		1147	+ dataCopyParams.srcStride = 0;
		1148	+ dataCopyParams.dstStride = 0;
		1149	+ DataCopyPadExtParams<bool> padParams;
		1150	+ padParams.isPad = true;
		1151	+ padParams.leftPadding = static_cast<uint8_t>(treeMaskStart % 32);
		1152	+ padParams.rightPadding = static_cast<uint8_t>(attenMaskSizeAlign - (attenMaskSize + treeMaskStart % 32));
		1153	+ padParams.paddingValue = 0;
		1154	+ DataCopyPad(attenMaskUb[(treeMaskStart - info.s2StartIdx) / 32 * 32],
		1155	+ srcGmAddr[maskOffset], dataCopyParams, padParams);
		1156	+ }
		1157	+ // else: curS2EndPos <= treeMaskStart，整个 tile 在零区，UB 已初始化为 0，无需拷贝
		1158	+ } else {
		1159	+ // 非 TREE：全量拷贝
		1160	+ uint64_t maskOffset = ComputeAttenMaskOffset<ENABLE_TREE>(info, 0, 0, isPre);
		1161	+ uint32_t attenMaskSizeAlign = Align(info.s2dealNum, 32U);
		1162	+ DataCopyExtParams dataCopyParams;
		1163	+ dataCopyParams.blockCount = 1;
		1164	+ dataCopyParams.blockLen = info.s2dealNum;
		1165	+ dataCopyParams.srcStride = info.attenMaskStride - info.s2dealNum;
		1166	+ dataCopyParams.dstStride = 0;
		1167	+ DataCopyPadExtParams<bool> padParams{true, 0,
		1168	+ static_cast<uint8_t>(attenMaskSizeAlign - info.s2dealNum), 0};
		1169	+ DataCopyPad(attenMaskUb, srcGmAddr[maskOffset], dataCopyParams, padParams);
		1170	+ }
1076		1171
1077	event_t enQueEvtID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));	1172	event_t enQueEvtID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
1078	SetFlag<HardEvent::MTE2_V>(enQueEvtID);	1173	SetFlag<HardEvent::MTE2_V>(enQueEvtID);
1079	WaitFlag<HardEvent::MTE2_V>(enQueEvtID);	1174	WaitFlag<HardEvent::MTE2_V>(enQueEvtID);
1080	for (uint32_t i = 1; i < info.gs1dealNum; i++) {	1175	for (uint32_t i = 1; i < info.gs1dealNum; i++) {
1081	- uint32_t offset = i * attenMaskSizeAlign;	1176	+ uint32_t offset = i * Align(info.s2dealNum, 32U);
1082	- DataCopy(attenMaskUb[offset], attenMaskUb, attenMaskSizeAlign);	1177	+ DataCopy(attenMaskUb[offset], attenMaskUb, Align(info.s2dealNum, 32U));
1083	}	1178	}
1084	}	1179	}
1085	}	1180	}

Mattention/fused_infer_attention_score/docs/aclnnFusedInferAttentionScoreV4.md+18-5

	@@ -210,6 +210,12 @@ aclnnStatus aclnnFusedInferAttentionScoreV4(
210	</ul>	210	</ul>
211	</li>	211	</li>
212	<li>sparseMode = 2、3、4时，attenMaskOptional的shape输入支持(2048, 2048)或(1,2048,2048)或(1,1,2048,2048)</li>	212	<li>sparseMode = 2、3、4时，attenMaskOptional的shape输入支持(2048, 2048)或(1,2048,2048)或(1,1,2048,2048)</li>
		213	+ <li>sparseMode = 9时：
		214	+ <ul>
		215	+ <li>inputLayout为BSH、BSND、BNSD时，attenMaskOptional的shape输入支持(B, Q_S, Q_S)</li>
		216	+ <li>inputLayout为TND时，attenMaskOptional的shape输入支持(∑Q_Si²,)，即每个batch的Q_Si×Q_Si mask拼接为1D tensor</li>
		217	+ </ul>
		218	+ </li>
213	</ul>	219	</ul>
214	</td>	220	</td>
215	<td>×</td>	221	<td>×</td>
	@@ -955,6 +961,11 @@ aclnnStatus aclnnFusedInferAttentionScoreV4(
955	<td>block_local</td>	961	<td>block_local</td>
956	<td>不支持</td>	962	<td>不支持</td>
957	</tr>	963	</tr>
		964	+ <tr>
		965	+ <td>9</td>
		966	+ <td>treeMask模式，用于推测解码场景的树形注意力掩码。需传入自定义的tree mask。</td>
		967	+ <td>非量化支持GQA和MLA场景，全量化仅支持MLA场景。不支持左padding、pseShift、sharedPrefix。输出dtype不支持INT8。每个batch需满足Q_S≤KV_S。inputLayout为BSH/BSND/BNSD时mask shape为(B,Q_S,Q_S)；inputLayout为TND时mask shape为(∑Q_Si²,)。</td>
		968	+ </tr>
958	</tbody>	969	</tbody>
959	</table>	970	</table>
960		971
	@@ -1492,7 +1503,7 @@ BFLOAT16和INT8不区分高精度和高性能，行无效修正对FLOAT16、BFLO
1492	</tr>	1503	</tr>
1493	<tr>	1504	<tr>
1494	<td>sparseMode</td>	1505	<td>sparseMode</td>
1495	- <td>支持0, 3, 4</td>	1506	+ <td>支持0, 3, 4, 9</td>
1496	</tr>	1507	</tr>
1497	<tr>	1508	<tr>
1498	<td rowspan="2">PagedAttention</td>	1509	<td rowspan="2">PagedAttention</td>
	@@ -1541,6 +1552,7 @@ BFLOAT16和INT8不区分高精度和高性能，行无效修正对FLOAT16、BFLO
1541	<li>sparse=0（attenMask为nullptr）</li>	1552	<li>sparse=0（attenMask为nullptr）</li>
1542	<li>sparse=3（传优化后的attenMask）</li>	1553	<li>sparse=3（传优化后的attenMask）</li>
1543	<li>sparse=4（传优化后的attenMask，需满足：Q_D=K_D=V_D≤256 或 Q_D=K_D=192且V_D=128/192；同时preTokens≥-actualSeqLengths、nextTokens≥-actualSeqLengthsKv、preTokens+nextTokens≥0）</li>	1554	<li>sparse=4（传优化后的attenMask，需满足：Q_D=K_D=V_D≤256 或 Q_D=K_D=192且V_D=128/192；同时preTokens≥-actualSeqLengths、nextTokens≥-actualSeqLengthsKv、preTokens+nextTokens≥0）</li>
		1555	+ <li>sparse=9（传入tree mask，inputLayout为BSH/BSND/BNSD时shape为(B,Q_S,Q_S)，inputLayout为TND时shape为(∑Q_Si²,)）</li>
1544	</ul>	1556	</ul>
1545	</li>	1557	</li>
1546	<li>innerPrecise：仅支持0（不带行无效的高精度模式）</li>	1558	<li>innerPrecise：仅支持0（不带行无效的高精度模式）</li>
	@@ -1641,7 +1653,7 @@ BFLOAT16和INT8不区分高精度和高性能，行无效修正对FLOAT16、BFLO
1641	<tr>	1653	<tr>
1642	<td>MASK</td>	1654	<td>MASK</td>
1643	<td>sparseMode</td>	1655	<td>sparseMode</td>
1644	- <td>sparseMode支持0, 3, 4</td>	1656	+ <td>sparseMode支持0, 3, 4, 9</td>
1645	<td>-</td>	1657	<td>-</td>
1646	</tr>	1658	</tr>
1647	<tr>	1659	<tr>
	@@ -1692,8 +1704,8 @@ BFLOAT16和INT8不区分高精度和高性能，行无效修正对FLOAT16、BFLO
1692	</tr>	1704	</tr>
1693	<tr>	1705	<tr>
1694	<td>sparseMode</td>	1706	<td>sparseMode</td>
1695	- <td>全量化场景sparseMode仅支持0,3</td>	1707	+ <td>全量化场景sparseMode仅支持0,3,9</td>
1696	- <td>qs=1时，仅支持sparseMode=0，且attenMask为nullptr; qs>1时，仅支持sparseMode=3，且attenMask的shape为[2048,2048]</td>	1708	+ <td>qs=1时，仅支持sparseMode=0，且attenMask为nullptr; qs>1时，支持sparseMode=3（attenMask的shape为[2048,2048]）或sparseMode=9（attenMask的shape见Mask章节）</td>
1697	</tr>	1709	</tr>
1698	<tr>	1710	<tr>
1699	<td>blockSize</td>	1711	<td>blockSize</td>
	@@ -1806,7 +1818,7 @@ BFLOAT16和INT8不区分高精度和高性能，行无效修正对FLOAT16、BFLO
1806	</tr>	1818	</tr>
1807	<tr>	1819	<tr>
1808	<td colspan="2">Mask</td>	1820	<td colspan="2">Mask</td>
1809	- <td colspan="3">当MTP等于0时，支持sparseMode=0且attenMask为nullptr；当MTP大于0、小于16时，支持sparseMode=3且传入优化后的attenMask矩阵，attenMask矩阵shape必须传入（2048*2048）；</td>	1821	+ <td colspan="3">当MTP等于0时，支持sparseMode=0且attenMask为nullptr；当MTP大于0、小于16时，支持sparseMode=3（传入优化后的attenMask矩阵，shape为2048*2048）或sparseMode=9（传入tree mask，inputLayout为BSH/BSND时shape为(B,Q_S,Q_S)，inputLayout为TND时shape为(∑Q_Si²,)）；</td>
1810	</tr>	1822	</tr>
1811	<tr>	1823	<tr>
1812	<td rowspan="9">伪量化</td>	1824	<td rowspan="9">伪量化</td>
	@@ -1977,6 +1989,7 @@ BFLOAT16和INT8不区分高精度和高性能，行无效修正对FLOAT16、BFLO
1977	<li>sparseMode = 0时，attenMaskOptional如果为空指针，或者在左padding场景传入attenMaskOptional，则忽略入参preTokens、nextTokens。</li>	1989	<li>sparseMode = 0时，attenMaskOptional如果为空指针，或者在左padding场景传入attenMaskOptional，则忽略入参preTokens、nextTokens。</li>
1978	<li>sparseMode = 2、3、4时，attenMaskOptional的shape需要为（2048,2048）或（1,2048,2048）或（1,1,2048,2048），且需要用户保证传入的attenMaskOptional为下三角，attenMaskOptional为nullptre或者传入的shape不正确报错。</li>	1990	<li>sparseMode = 2、3、4时，attenMaskOptional的shape需要为（2048,2048）或（1,2048,2048）或（1,1,2048,2048），且需要用户保证传入的attenMaskOptional为下三角，attenMaskOptional为nullptre或者传入的shape不正确报错。</li>
1979	<li>sparseMode = 1、2、3的场景忽略入参preTokens、nextTokens并按照相关规则赋值。</li>	1991	<li>sparseMode = 1、2、3的场景忽略入参preTokens、nextTokens并按照相关规则赋值。</li>
		1992	+ <li>sparseMode = 9时，非量化支持GQA和MLA场景，全量化仅支持MLA场景。attenMaskOptional不能为空。inputLayout为BSH/BSND/BNSD时shape为(B,Q_S,Q_S)；inputLayout为TND时shape为(∑Q_Si²,)。不支持左padding、pseShift、sharedPrefix。输出dtype不支持INT8。每个batch需满足Q_S≤KV_S。</li>
1980	<li>sparseMode取其它值时会报错</li>	1993	<li>sparseMode取其它值时会报错</li>
1981	</ul>	1994	</ul>
1982	</td>	1995	</td>

Mattention/fused_infer_attention_score/op_host/arch32/fused_infer_attention_score_tiling_check.h+1-0文件内容审核中，请稍后刷新重试

Mattention/fused_infer_attention_score/op_host/arch32/fused_infer_attention_score_tiling_check_consistency.cpp+31-1文件内容审核中，请稍后刷新重试

Mattention/fused_infer_attention_score/op_host/arch32/fused_infer_attention_score_tiling_check_feature.cpp+68-4

	@@ -160,8 +160,71 @@ ge::graphStatus FiaTilingCheck::CheckFeatureMla()
160	return ge::GRAPH_SUCCESS;	160	return ge::GRAPH_SUCCESS;
161	}	161	}
162		162
		163	+// 补充sparse9特性的拦截校验
		164	+ge::graphStatus FiaTilingCheck::CheckFeatureSparseMode() const
		165	+{
		166	+ int32_t sparseMode = fiaInfo_.sparseMode;
		167	+ // sparse9 仅在rope分离场景下存在，不支持左padding、PSE、公共前缀、后量化等特性拦截s2 >= s1
		168	+ if (sparseMode == SPARSE_MODE_TREE) {
		169	+ // 特性校验
		170	+ OP_CHECK_IF(ropeMode_ != RopeMode::ROPE_SPLIT,
		171	+ OP_LOGE(opName_,
		172	+ "In %s situation, when query_rope and key_rope not exist, %s does not support sparse(%d).",
		173	+ QuantModeToSerialString(quantMode_).c_str(), SPARSE_MODE_NAME.c_str(), sparseMode),
		174	+ return ge::GRAPH_FAILED);
		175	+
		176	+ OP_CHECK_IF(fiaInfo_.qPaddingSizeFlag \|\| fiaInfo_.kvPaddingSizeFlag,
		177	+ OP_LOGE(opName_,
		178	+ "In %s situation, when sparse is %d, query_padding_size or kv_padding_size should be not exist.",
		179	+ QuantModeToSerialString(quantMode_).c_str(), sparseMode),
		180	+ return ge::GRAPH_FAILED);
		181	+
		182	+ OP_CHECK_IF(fiaInfo_.pseShiftFlag,
		183	+ OP_LOGE(opName_,
		184	+ "In %s situation, when sparse is %d, pse_shift should be not exist.",
		185	+ QuantModeToSerialString(quantMode_).c_str(), sparseMode),
		186	+ return ge::GRAPH_FAILED);
		187	+
		188	+ OP_CHECK_IF(fiaInfo_.sysPrefixFlag,
		189	+ OP_LOGE(opName_,
		190	+ "In %s situation, when sparse is %d, key_shared_prefix and key_shared_prefix should be not exist.",
		191	+ QuantModeToSerialString(quantMode_).c_str(), sparseMode),
		192	+ return ge::GRAPH_FAILED);
		193	+
		194	+ OP_CHECK_IF(outputType_ == ge::DT_INT8,
		195	+ OP_LOGE(opName_,
		196	+ "In %s situation, when sparse is %d, output dtype %d is not currently supported.",
		197	+ QuantModeToSerialString(quantMode_).c_str(), sparseMode, static_cast<int32_t>(outputType_)),
		198	+ return ge::GRAPH_FAILED);
		199	+
		200	+ // s2 >= s1拦截
		201	+ // tiling下沉场景由于actualSeqlen得不到，所以不进行校验
		202	+ if (fiaInfo_.isMaxWorkspace) {
		203	+ return ge::GRAPH_SUCCESS;
		204	+ }
		205	+ // s2=0 的 batch（入图padding或空tensor场景）不校验 s1<=s2
		206	+ int32_t actualSeqSize = std::min(qSize.size(), kvSize.size());
		207	+ for (int32_t i = 0; i < actualSeqSize; i++) {
		208	+ if (kvSize[i] == 0) {
		209	+ continue;
		210	+ }
		211	+ OP_CHECK_IF(qSize[i] > kvSize[i],
		212	+ OP_LOGE(opName_,
		213	+ "In %s situation, when sparse is %d, qSize[%d] should less than or equal to kvSize[%d],"
		214	+ "but got qSize %d and kvSize %d.",
		215	+ QuantModeToSerialString(quantMode_).c_str(), sparseMode, i, i, qSize[i], kvSize[i]),
		216	+ return ge::GRAPH_FAILED);
		217	+ }
		218	+ }
		219	+ return ge::GRAPH_SUCCESS;
		220	+}
		221	+
163	ge::graphStatus FiaTilingCheck::CheckFeatureMask() const	222	ge::graphStatus FiaTilingCheck::CheckFeatureMask() const
164	{	223	{
		224	+ if (ge::GRAPH_SUCCESS != CheckFeatureSparseMode()) {
		225	+ return ge::GRAPH_FAILED;
		226	+ }
		227	+
165	if ((!attenMaskFlag_) && (fiaInfo_.sparseMode != SPARSE_MODE_NO_MASK)) {	228	if ((!attenMaskFlag_) && (fiaInfo_.sparseMode != SPARSE_MODE_NO_MASK)) {
166	OP_LOGE(opName_, "when %s is %d, it not 0, %s should not be null.",	229	OP_LOGE(opName_, "when %s is %d, it not 0, %s should not be null.",
167	SPARSE_MODE_NAME.c_str(), fiaInfo_.sparseMode, ATTEN_MASK_NAME.c_str());	230	SPARSE_MODE_NAME.c_str(), fiaInfo_.sparseMode, ATTEN_MASK_NAME.c_str());
	@@ -194,10 +257,11 @@ ge::graphStatus FiaTilingCheck::CheckFeatureMask() const
194		257
195	if (ropeMode_ == RopeMode::ROPE_SPLIT && vHeadDim_ == 512U) {	258	if (ropeMode_ == RopeMode::ROPE_SPLIT && vHeadDim_ == 512U) {
196	int32_t sparseMode = fiaInfo_.sparseMode;	259	int32_t sparseMode = fiaInfo_.sparseMode;
197	- if (sparseMode != SPARSE_MODE_NO_MASK && sparseMode != SPARSE_MODE_RIGHT_DOWN && sparseMode != SPARSE_MODE_BAND) {	260	+ if (sparseMode != SPARSE_MODE_NO_MASK && sparseMode != SPARSE_MODE_RIGHT_DOWN &&
		261	+ sparseMode != SPARSE_MODE_BAND && sparseMode != SPARSE_MODE_TREE) {
198	OP_LOGE(opName_,	262	OP_LOGE(opName_,
199	"In %s situation, when query_rope and key_rope exsists and the head dim of value is %u, %s only "	263	"In %s situation, when query_rope and key_rope exsists and the head dim of value is %u, %s only "
200	- "support 0/3/4, but got %d.",	264	+ "support 0/3/4/9, but got %d.",
201	QuantModeToSerialString(quantMode_).c_str(), vHeadDim_, SPARSE_MODE_NAME.c_str(), sparseMode);	265	QuantModeToSerialString(quantMode_).c_str(), vHeadDim_, SPARSE_MODE_NAME.c_str(), sparseMode);
202	return ge::GRAPH_FAILED;	266	return ge::GRAPH_FAILED;
203	}	267	}
	@@ -229,13 +293,13 @@ ge::graphStatus FiaTilingCheck::CheckFeatureLeftPadding() const
229		293
230	OP_CHECK_IF(kvStorageMode_ == KvStorageMode::TENSOR_LIST,	294	OP_CHECK_IF(kvStorageMode_ == KvStorageMode::TENSOR_LIST,
231	OP_LOGE(opName_,	295	OP_LOGE(opName_,
232	- "when query_padding_size or kv_padding_size exists, key/value tensorlist is not suppoprted; in this "	296	+ "when query_padding_size or kv_padding_size exists, key/value tensorlist is not supported; in this "
233	"case, the tensor number of key/value should be 1"),	297	"case, the tensor number of key/value should be 1"),
234	return ge::GRAPH_FAILED);	298	return ge::GRAPH_FAILED);
235		299
236	OP_CHECK_IF(kvStorageMode_ == KvStorageMode::PAGE_ATTENTION,	300	OP_CHECK_IF(kvStorageMode_ == KvStorageMode::PAGE_ATTENTION,
237	OP_LOGE(opName_,	301	OP_LOGE(opName_,
238	- "when query_padding_size or kv_padding_size exists, page attention is not suppoprted; in this case, "	302	+ "when query_padding_size or kv_padding_size exists, page attention is not supported; in this case, "
239	"block_table should exist and block_size is not 0"),	303	"block_table should exist and block_size is not 0"),
240	return ge::GRAPH_FAILED);	304	return ge::GRAPH_FAILED);
241		305

Mattention/fused_infer_attention_score/op_host/arch32/fused_infer_attention_score_tiling_check_single_para.cpp+3-2

	@@ -633,10 +633,11 @@ ge::graphStatus FiaTilingCheck::CheckSingleParaSparseMode() const
633	SPARSE_MODE_ALL_MASK,	633	SPARSE_MODE_ALL_MASK,
634	SPARSE_MODE_LEFT_UP,	634	SPARSE_MODE_LEFT_UP,
635	SPARSE_MODE_RIGHT_DOWN,	635	SPARSE_MODE_RIGHT_DOWN,
636	- SPARSE_MODE_BAND	636	+ SPARSE_MODE_BAND,
		637	+ SPARSE_MODE_TREE
637	};	638	};
638	if (ge::GRAPH_SUCCESS != CheckAttrValueSupport(opParamInfo_.sparseMode, sparseModeList, SPARSE_MODE_NAME)) {	639	if (ge::GRAPH_SUCCESS != CheckAttrValueSupport(opParamInfo_.sparseMode, sparseModeList, SPARSE_MODE_NAME)) {
639	- OP_LOGE(opName_, "sparseMode only supports 0/1/2/3/4, but got %u", *opParamInfo_.sparseMode);	640	+ OP_LOGE(opName_, "sparseMode only supports 0/1/2/3/4/9, but got %u", *opParamInfo_.sparseMode);
640	return ge::GRAPH_FAILED;	641	return ge::GRAPH_FAILED;
641	}	642	}
642	return ge::GRAPH_SUCCESS;	643	return ge::GRAPH_SUCCESS;

Mattention/fused_infer_attention_score/op_host/arch32/fused_infer_attention_score_tiling_info_parser.cpp+32-0文件内容审核中，请稍后刷新重试

Mattention/fused_infer_attention_score/op_host/arch32/fused_infer_attention_score_tiling_info_parser.h+1-0文件内容审核中，请稍后刷新重试

Mattention/fused_infer_attention_score/op_host/arch35/fused_infer_attention_score_tiling_v2.cpp+9-0文件内容审核中，请稍后刷新重试

Mattention/fused_infer_attention_score/op_host/fused_infer_attention_score_tiling.cpp+5-2文件内容审核中，请稍后刷新重试

Mattention/incre_flash_attention/op_host/incre_flash_attention_tiling.cpp+10-7

	@@ -537,7 +537,7 @@ ge::graphStatus IFATiling::GetRopeAndGqaFlag(const uint32_t sOfQuery, const uint
537	OP_LOGE(ifaContext_->opName, "In case where MLA is not applied, S of Query:%u is invalid. It should be in range [1, 16]", sOfQuery),	537	OP_LOGE(ifaContext_->opName, "In case where MLA is not applied, S of Query:%u is invalid. It should be in range [1, 16]", sOfQuery),
538	return ge::GRAPH_FAILED);	538	return ge::GRAPH_FAILED);
539	} else if (layout != "TND" && layout != "TND_NTD") {	539	} else if (layout != "TND" && layout != "TND_NTD") {
540	- OP_CHECK_IF(sOfQuery > 16, OP_LOGE(ifaContext_->opName, "QueryS(%u) should not be bigger than 16 in MLA.", sOfQuery),	540	+ OP_CHECK_IF(sOfQuery > 32, OP_LOGE(ifaContext_->opName, "QueryS(%u) should not be bigger than 32 in MLA.", sOfQuery),
541	return ge::GRAPH_FAILED);	541	return ge::GRAPH_FAILED);
542	}	542	}
543	OP_CHECK_IF(layout == "TND" && headDim_ == 512 && !ropeFlag_, OP_LOGE(ifaContext_->opName,	543	OP_CHECK_IF(layout == "TND" && headDim_ == 512 && !ropeFlag_, OP_LOGE(ifaContext_->opName,
	@@ -588,12 +588,13 @@ ge::graphStatus IFATiling::QKVPreProcess4TND(const std::string layout)
588	std::vector<int64_t> actualSeqQ(actualLenQDims_);	588	std::vector<int64_t> actualSeqQ(actualLenQDims_);
589	int64_t tmpQSeqSize = 0;	589	int64_t tmpQSeqSize = 0;
590		590
591	- for (int b = 0; b < static_cast<int>(actualLenQDims_); b++) {	591	+ for (int32_t b = 0; b < static_cast<int>(actualLenQDims_); b++) {
592	actualSeqQ[b] = (b <= 0) ? actualSeqQTnd[0] : (actualSeqQTnd[b] - actualSeqQTnd[b - 1]);	592	actualSeqQ[b] = (b <= 0) ? actualSeqQTnd[0] : (actualSeqQTnd[b] - actualSeqQTnd[b - 1]);
593	- OP_CHECK_IF((actualSeqQ[b] < 0) \|\| (actualSeqQ[b] > 16), // 16 MTP最大QS	593	+ OP_CHECK_IF((actualSeqQ[b] < 0) \|\| (actualSeqQ[b] > 32), // 32 MTP最大QS
594	- OP_LOGE(ifaContext_->opName, "%s QS(%ld) of batch(%d) computed by the query's actual sequence lengths should be in range [0, 16].", layout.c_str(), actualSeqQ[b], b),	594	+ OP_LOGE(ifaContext_->opName, "%s QS(%ld) of batch(%d) computed by the query's actual sequence lengths should be in range [0, 32].", layout.c_str(), actualSeqQ[b], b),
595	return ge::GRAPH_FAILED);	595	return ge::GRAPH_FAILED);
596	tmpQSeqSize = std::max(tmpQSeqSize, actualSeqQ[b]);	596	tmpQSeqSize = std::max(tmpQSeqSize, actualSeqQ[b]);
		597	+ qSeqSquareSum_ += actualSeqQ[b] * actualSeqQ[b];
597	}	598	}
598		599
599	OP_CHECK_IF((tSeqSize_ != actualSeqQTnd[actualLenQDims_ - 1]),	600	OP_CHECK_IF((tSeqSize_ != actualSeqQTnd[actualLenQDims_ - 1]),
	@@ -1973,14 +1974,14 @@ void IFATiling::GetActualSeqInfo(const int64_t *actualSeqKv, ActualSeqInfo &actu
1973	// TND格式，actual_seq_q定义为累积长度，这里做转化再分核	1974	// TND格式，actual_seq_q定义为累积长度，这里做转化再分核
1974	const int64_t *actualSeqQTnd = ifaContext_->actualSeqLengthsQ.tensor->GetData<int64_t>();	1975	const int64_t *actualSeqQTnd = ifaContext_->actualSeqLengthsQ.tensor->GetData<int64_t>();
1975	actualSeqInfo.actualSeqQ[0] = actualSeqQTnd[0];	1976	actualSeqInfo.actualSeqQ[0] = actualSeqQTnd[0];
1976	- for (int b = 1; b < static_cast<int>(bSize); b++) {	1977	+ for (int32_t b = 1; b < static_cast<int>(bSize); b++) {
1977	actualSeqInfo.actualSeqQ[b] = actualSeqQTnd[b] - actualSeqQTnd[b - 1];	1978	actualSeqInfo.actualSeqQ[b] = actualSeqQTnd[b] - actualSeqQTnd[b - 1];
1978	if (actualLenDims_ != 1U) {	1979	if (actualLenDims_ != 1U) {
1979	actualSeqInfo.maxActualseqkv = std::max(actualSeqInfo.maxActualseqkv, actualSeqKv[b]);	1980	actualSeqInfo.maxActualseqkv = std::max(actualSeqInfo.maxActualseqkv, actualSeqKv[b]);
1980	}	1981	}
1981	}	1982	}
1982	} else {	1983	} else {
1983	- for (int b = 0; b < static_cast<int>(bSize); b++) {	1984	+ for (int32_t b = 0; b < static_cast<int>(bSize); b++) {
1984	actualSeqInfo.actualSeqQ[b] = qSeqSize_; // 需要检查	1985	actualSeqInfo.actualSeqQ[b] = qSeqSize_; // 需要检查
1985	if (actualLenDims_ != 1U) {	1986	if (actualLenDims_ != 1U) {
1986	actualSeqInfo.maxActualseqkv = std::max(actualSeqInfo.maxActualseqkv, actualSeqKv[b]);	1987	actualSeqInfo.maxActualseqkv = std::max(actualSeqInfo.maxActualseqkv, actualSeqKv[b]);
	@@ -3220,8 +3221,10 @@ void IFATiling::FillTilingBaseParamsMla()
3220	tilingDataMla_.baseParams.set_actualLenDims(actualLenDims_);	3221	tilingDataMla_.baseParams.set_actualLenDims(actualLenDims_);
3221	tilingDataMla_.baseParams.set_attenMaskFlag(attenMaskFlag_ ? 1 : 0);	3222	tilingDataMla_.baseParams.set_attenMaskFlag(attenMaskFlag_ ? 1 : 0);
3222	tilingDataMla_.baseParams.set_attenMaskSize(attenMaskSize_);	3223	tilingDataMla_.baseParams.set_attenMaskSize(attenMaskSize_);
		3224	+ tilingDataMla_.baseParams.set_sparseMode(sparseMode_);
3223	tilingDataMla_.baseParams.set_outputLayout(static_cast<uint32_t>(outputLayout_));	3225	tilingDataMla_.baseParams.set_outputLayout(static_cast<uint32_t>(outputLayout_));
3224	tilingDataMla_.baseParams.set_softmaxLseFlag(softmaxLseFlag_ ? 1 : 0);	3226	tilingDataMla_.baseParams.set_softmaxLseFlag(softmaxLseFlag_ ? 1 : 0);
		3227	+ tilingDataMla_.baseParams.set_tSeqSize(tSeqSize_);
3225	}	3228	}
3226		3229
3227	// for flash decode	3230	// for flash decode
	@@ -4107,7 +4110,7 @@ uint32_t IFATiling::GetTotalQBlockNum() const
4107	uint32_t totalQblockSum = 0;	4110	uint32_t totalQblockSum = 0;
4108	uint32_t curSeqLenQ = 0;	4111	uint32_t curSeqLenQ = 0;
4109	uint32_t preSeqLenQ = 0;	4112	uint32_t preSeqLenQ = 0;
4110	- for (int bIdx = 0; bIdx < static_cast<int>(actualLenQDims_); bIdx++) {	4113	+ for (int32_t bIdx = 0; bIdx < static_cast<int>(actualLenQDims_); bIdx++) {
4111	// actualLenDataQ里的值单调递增	4114	// actualLenDataQ里的值单调递增
4112	curSeqLenQ = static_cast<uint32_t>(actualLenDataQ[bIdx]);	4115	curSeqLenQ = static_cast<uint32_t>(actualLenDataQ[bIdx]);
4113	uint32_t tmpBlkNum = static_cast<uint32_t>(curSeqLenQ - preSeqLenQ + seqStepQ_ - 1) / seqStepQ_;	4116	uint32_t tmpBlkNum = static_cast<uint32_t>(curSeqLenQ - preSeqLenQ + seqStepQ_ - 1) / seqStepQ_;

Mattention/incre_flash_attention/op_host/incre_flash_attention_tiling_check.cpp+112-8

	@@ -329,7 +329,7 @@ ge::graphStatus IFATiling::CheckInputQKVTypeMatch() const
329		329
330	OP_CHECK_IF(	330	OP_CHECK_IF(
331	((!ropeFlag_) && inputQType_ == ge::DT_INT8 && inputKvType_ == ge::DT_INT8),	331	((!ropeFlag_) && inputQType_ == ge::DT_INT8 && inputKvType_ == ge::DT_INT8),
332	- OP_LOGE(ifaContext_->opName, "When QueryRope/KeyRope is null and Qs(%u) in [1, 16], not support qkv datatype all int8.", qSeqSize_), return ge::GRAPH_FAILED);	332	+ OP_LOGE(ifaContext_->opName, "When QueryRope/KeyRope is null and Qs(%u) in [1, 32], not support qkv datatype all int8.", qSeqSize_), return ge::GRAPH_FAILED);
333		333
334	OP_CHECK_IF(	334	OP_CHECK_IF(
335	((inputQType_ == ge::DT_FLOAT16) && (inputKvType_ != ge::DT_FLOAT16 && inputKvType_ != ge::DT_INT8 && inputKvType_ != ge::DT_INT4)), OP_LOGE(ifaContext_->opName, "when input Q type is fp16, KV type %d should be fp16 or int8 or int4", inputKvType_),	335	((inputQType_ == ge::DT_FLOAT16) && (inputKvType_ != ge::DT_FLOAT16 && inputKvType_ != ge::DT_INT8 && inputKvType_ != ge::DT_INT4)), OP_LOGE(ifaContext_->opName, "when input Q type is fp16, KV type %d should be fp16 or int8 or int4", inputKvType_),
	@@ -1345,6 +1345,87 @@ ge::graphStatus IFATiling::ProcessPseShift()
1345	return ge::GRAPH_SUCCESS;	1345	return ge::GRAPH_SUCCESS;
1346	}	1346	}
1347		1347
		1348	+ge::graphStatus IFATiling::CheckTreeSparseMask()
		1349	+{
		1350	+ // 仅支持MLA全量化，且要求s2 >= s1
		1351	+ OP_CHECK_IF(socVersion_ != IfaSocVersion::SOC_ASCEND_910B,
		1352	+ OP_LOGE(ifaContext_->opName, "Tree sparse9 is only supported in ASCEND_910B."),
		1353	+ return ge::GRAPH_FAILED);
		1354	+
		1355	+ OP_CHECK_IF(!(ropeFlag_ && quantFlag_),
		1356	+ OP_LOGE(ifaContext_->opName, "Tree Sparse9 is only supported in MLA full quant situation."),
		1357	+ return ge::GRAPH_FAILED);
		1358	+
		1359	+
		1360	+ // tiling下沉场景，获取不到actualseqlen,不进行校验
		1361	+ if (isWorkspace_) {
		1362	+ return ge::GRAPH_SUCCESS;
		1363	+ }
		1364	+
		1365	+ // 补充s2 >= s1的拦截，在MLA全量化场景，必须开启PA(NZ)，所以actualSeqKV不进行累加
		1366	+ if (inputLayout_ == IfaLayout::TND) {
		1367	+ const int64_t *actualSeqQTnd = ifaContext_->actualSeqLengthsQ.tensor->GetData<int64_t>();
		1368	+ const int64_t *actualSeqKVTnd = ifaContext_->actualSeqLengths.tensor->GetData<int64_t>();
		1369	+ int64_t qActSize = 0;
		1370	+ int64_t kvActSize = 0;
		1371	+
		1372	+ // s2=0 的 batch（入图padding或空tensor场景）不校验 s1<=s2
		1373	+ for (int32_t b = 0; b < actualLenQDims_; b++) {
		1374	+ kvActSize = actualSeqKVTnd[b];
		1375	+ if (kvActSize == 0) {
		1376	+ continue;
		1377	+ }
		1378	+ qActSize = (b == 0) ? actualSeqQTnd[0] : (actualSeqQTnd[b] - actualSeqQTnd[b - 1]);
		1379	+ OP_CHECK_IF(qActSize > kvActSize,
		1380	+ OP_LOGE(ifaContext_->opName,
		1381	+ "In MLA full quant situation, when sparse is %d, qSize(%ld) should less than or equal to kvSize(%ld).",
		1382	+ sparseMode_, qActSize, kvActSize),
		1383	+ return ge::GRAPH_FAILED);
		1384	+ }
		1385	+ } else {
		1386	+ OP_CHECK_IF(qSeqSize_ > seqSize_,
		1387	+ OP_LOGE(ifaContext_->opName,
		1388	+ "In MLA full quant situation, when sparse is %d, qSize(%ld) should less than or equal to kvSize(%ld).",
		1389	+ sparseMode_, qSeqSize_, seqSize_),
		1390	+ return ge::GRAPH_FAILED);
		1391	+ }
		1392	+ return ge::GRAPH_SUCCESS;
		1393	+}
		1394	+
		1395	+ge::graphStatus IFATiling::CheckTreeSparseMaskShape()
		1396	+{
		1397	+ // sparse9时需要传入Mask，在TND场景，传入∑s1²，一维矩阵
		1398	+ // 非TND场景传入[B, S1, S1]
		1399	+ auto maskShape = ifaContext_->attenMask.tensor;
		1400	+ if (inputLayout_ == IfaLayout::TND) {
		1401	+ // sparse9时需要传入Mask，在TND场景，传入∑s1²，一维矩阵
		1402	+ OP_CHECK_IF(maskShape == nullptr,
		1403	+ OP_LOGE(ifaContext_->opName, "TND/TND_NTD need input attenMask when sparse = 9."), return ge::GRAPH_FAILED);
		1404	+
		1405	+ auto shape = ifaContext_->attenMask.tensor->GetStorageShape();
		1406	+ OP_CHECK_IF(shape.GetDimNum() != 1U,
		1407	+ OP_LOGE(ifaContext_->opName, "TND/TND_NTD when sparse = 9, atten_mask tensor shape must be 1 dimension."),
		1408	+ return ge::GRAPH_FAILED);
		1409	+
		1410	+ if (!isWorkspace_) {
		1411	+ OP_CHECK_IF(shape.GetDim(0) != qSeqSquareSum_,
		1412	+ OP_LOGE(ifaContext_->opName, "TND/TND_NTD when sparse = 9, atten_mask tensor shape must be ∑s1²."),
		1413	+ return ge::GRAPH_FAILED);
		1414	+ }
		1415	+ attenMaskFlag_ = true;
		1416	+ } else {
		1417	+ OP_CHECK_IF(maskShape == nullptr,
		1418	+ OP_LOGE(ifaContext_->opName, "When sparse = 9, atten_mask should not be null."), return ge::GRAPH_FAILED);
		1419	+
		1420	+ auto shape = ifaContext_->attenMask.tensor->GetStorageShape();
		1421	+ OP_CHECK_IF(shape.GetDimNum() != 3U \|\| shape.GetDim(0) != batchSize_ \|\| shape.GetDim(1) != qSeqSize_ \|\| shape.GetDim(2) != qSeqSize_ ,
		1422	+ OP_LOGE(ifaContext_->opName, "BSH when sparse = 9, atten_mask tensor shape must be [B, S1, S1]."),
		1423	+ return ge::GRAPH_FAILED);
		1424	+ attenMaskFlag_ = true;
		1425	+ }
		1426	+ return ge::GRAPH_SUCCESS;
		1427	+}
		1428	+
1348	ge::graphStatus IFATiling::CheckTndMaskShapeWithSparseMode()	1429	ge::graphStatus IFATiling::CheckTndMaskShapeWithSparseMode()
1349	{	1430	{
1350	auto maskShape = ifaContext_->attenMask.tensor; // input shape = 4	1431	auto maskShape = ifaContext_->attenMask.tensor; // input shape = 4
	@@ -1361,8 +1442,13 @@ ge::graphStatus IFATiling::CheckTndMaskShapeWithSparseMode()
1361	OP_LOGE(ifaContext_->opName, "TND/TND_NTD when sparse = 3, atten_mask tensor shape must be [2048, 2048]."),	1442	OP_LOGE(ifaContext_->opName, "TND/TND_NTD when sparse = 3, atten_mask tensor shape must be [2048, 2048]."),
1362	return ge::GRAPH_FAILED);	1443	return ge::GRAPH_FAILED);
1363	attenMaskFlag_ = true;	1444	attenMaskFlag_ = true;
		1445	+ } else if (sparseMode_ == 9U) {
		1446	+ if (ge::GRAPH_SUCCESS == CheckTreeSparseMask() &&
		1447	+ ge::GRAPH_SUCCESS == CheckTreeSparseMaskShape()) {
		1448	+ attenMaskFlag_ = true;
		1449	+ }
1364	} else {	1450	} else {
1365	- OP_LOGE(ifaContext_->opName, "TND/TND_NTD only support sparse(%u) = 0 or 3.", sparseMode_);	1451	+ OP_LOGE(ifaContext_->opName, "TND/TND_NTD only support sparse(%u) = 0/3/9.", sparseMode_);
1366	return ge::GRAPH_FAILED;	1452	return ge::GRAPH_FAILED;
1367	}	1453	}
1368	return ge::GRAPH_SUCCESS;	1454	return ge::GRAPH_SUCCESS;
	@@ -1370,14 +1456,23 @@ ge::graphStatus IFATiling::CheckTndMaskShapeWithSparseMode()
1370		1456
1371	ge::graphStatus IFATiling::CheckMaskShapeWithQSeq() const	1457	ge::graphStatus IFATiling::CheckMaskShapeWithQSeq() const
1372	{	1458	{
1373	- if (antiQuantFlag_ \|\| (ifaContext_->dequantScaleQuery.tensor != nullptr && ropeFlag_)) {	1459	+ if (quantFlag_) {
1374	- OP_CHECK_IF((ropeFlag_ && qSeqSize_ > 1U && static_cast<int32_t>(sparseMode_) != 3),	1460	+ OP_CHECK_IF((ropeFlag_ && qSeqSize_ > 1U && (static_cast<int32_t>(sparseMode_) != 3 && static_cast<int32_t>(sparseMode_) != 9)),
1375	- OP_LOGE(ifaContext_->opName, "when queryS > 1, sparseMode(%d) only support 3 "	1461	+ OP_LOGE(ifaContext_->opName, "when queryS > 1, sparseMode(%d) only support 3/9 "
1376	- "in MLA when antiquant or full quant situation.", static_cast<int32_t>(sparseMode_)),	1462	+ "in MLA when full quant situation.", static_cast<int32_t>(sparseMode_)),
1377	- return ge::GRAPH_FAILED);	1463	+ return ge::GRAPH_FAILED);
1378	OP_CHECK_IF((ropeFlag_ && qSeqSize_ == 1U && static_cast<int32_t>(sparseMode_) != 0),	1464	OP_CHECK_IF((ropeFlag_ && qSeqSize_ == 1U && static_cast<int32_t>(sparseMode_) != 0),
1379	OP_LOGE(ifaContext_->opName, "when queryS = 1, sparseMode(%d) only support 0 "	1465	OP_LOGE(ifaContext_->opName, "when queryS = 1, sparseMode(%d) only support 0 "
1380	- "in MLA when antiquant or full quant situation.", static_cast<int32_t>(sparseMode_)),	1466	+ "in MLA when full quant situation.", static_cast<int32_t>(sparseMode_)),
		1467	+ return ge::GRAPH_FAILED);
		1468	+ } else if (antiQuantFlag_) {
		1469	+ OP_CHECK_IF((ropeFlag_ && qSeqSize_ > 1U && static_cast<int32_t>(sparseMode_) != 3),
		1470	+ OP_LOGE(ifaContext_->opName, "when queryS > 1, sparseMode(%d) only support 3 "
		1471	+ "in MLA when antiquant situation.", static_cast<int32_t>(sparseMode_)),
		1472	+ return ge::GRAPH_FAILED);
		1473	+ OP_CHECK_IF((ropeFlag_ && qSeqSize_ == 1U && static_cast<int32_t>(sparseMode_) != 0),
		1474	+ OP_LOGE(ifaContext_->opName, "when queryS = 1, sparseMode(%d) only support 0 "
		1475	+ "in MLA when antiquant situation.", static_cast<int32_t>(sparseMode_)),
1381	return ge::GRAPH_FAILED);	1476	return ge::GRAPH_FAILED);
1382	} else {	1477	} else {
1383	OP_CHECK_IF((ropeFlag_ && qSeqSize_ > 1U &&	1478	OP_CHECK_IF((ropeFlag_ && qSeqSize_ > 1U &&
	@@ -1414,6 +1509,11 @@ ge::graphStatus IFATiling::CheckAttenMaskShape()
1414	OP_LOGW(ifaContext_->opName, "atten_mask tensor exist, but atten_mask shape size is 0.");	1509	OP_LOGW(ifaContext_->opName, "atten_mask tensor exist, but atten_mask shape size is 0.");
1415	return ge::GRAPH_SUCCESS;	1510	return ge::GRAPH_SUCCESS;
1416	}	1511	}
		1512	+
		1513	+ if (sparseMode_ == 9U && CheckTreeSparseMaskShape() == ge::GRAPH_SUCCESS) {
		1514	+ attenMaskFlag_ = true;
		1515	+ return ge::GRAPH_SUCCESS;
		1516	+ }
1417	return ge::GRAPH_FAILED;	1517	return ge::GRAPH_FAILED;
1418	}	1518	}
1419		1519
	@@ -1430,6 +1530,10 @@ ge::graphStatus IFATiling::ProcessAttenMask()
1430	return ge::GRAPH_FAILED;	1530	return ge::GRAPH_FAILED;
1431	}	1531	}
1432		1532
		1533	+ if (sparseMode_ == 9U && CheckTreeSparseMask() != ge::GRAPH_SUCCESS) {
		1534	+ return ge::GRAPH_FAILED;
		1535	+ }
		1536	+
1433	if (CheckAttenMaskShape() == ge::GRAPH_SUCCESS) {	1537	if (CheckAttenMaskShape() == ge::GRAPH_SUCCESS) {
1434	return ge::GRAPH_SUCCESS;	1538	return ge::GRAPH_SUCCESS;
1435	}	1539	}

Mattention/incre_flash_attention/op_host/incre_flash_attention_tiling_impl.h+3-0

	@@ -85,6 +85,8 @@ private:
85	ge::graphStatus ProcessBaseInputs();	85	ge::graphStatus ProcessBaseInputs();
86	ge::graphStatus ProcessOptionalTensors();	86	ge::graphStatus ProcessOptionalTensors();
87	ge::graphStatus ProcessPseShift();	87	ge::graphStatus ProcessPseShift();
		88	+ ge::graphStatus CheckTreeSparseMask();
		89	+ ge::graphStatus CheckTreeSparseMaskShape();
88	ge::graphStatus CheckTndMaskShapeWithSparseMode();	90	ge::graphStatus CheckTndMaskShapeWithSparseMode();
89	ge::graphStatus CheckMaskShapeWithQSeq() const;	91	ge::graphStatus CheckMaskShapeWithQSeq() const;
90	ge::graphStatus CheckAttenMaskShape();	92	ge::graphStatus CheckAttenMaskShape();
	@@ -360,6 +362,7 @@ private:
360	uint32_t sMax_ = 0;	362	uint32_t sMax_ = 0;
361	uint32_t tSeqSize_ = 1; // TND格式T轴长度	363	uint32_t tSeqSize_ = 1; // TND格式T轴长度
362	uint32_t qSeqSize_ = 1; // 默认S1 = 1	364	uint32_t qSeqSize_ = 1; // 默认S1 = 1
		365	+ uint64_t qSeqSquareSum_ = 0; // 给Sparse9 TND下mask大小用
363	uint32_t blockTypeSize_ = 0; // 计算中间量大小	366	uint32_t blockTypeSize_ = 0; // 计算中间量大小
364	uint32_t kvSplitPart_ = 1;	367	uint32_t kvSplitPart_ = 1;
365		368

Mattention/incre_flash_attention/op_host/incre_flash_attention_tiling_mla.h+2-0

	@@ -37,8 +37,10 @@ TILING_DATA_FIELD_DEF(uint32_t, actualLenDims)
37	TILING_DATA_FIELD_DEF(uint32_t, antiquantMode)	37	TILING_DATA_FIELD_DEF(uint32_t, antiquantMode)
38	TILING_DATA_FIELD_DEF(uint32_t, attenMaskFlag)	38	TILING_DATA_FIELD_DEF(uint32_t, attenMaskFlag)
39	TILING_DATA_FIELD_DEF(uint32_t, attenMaskSize)	39	TILING_DATA_FIELD_DEF(uint32_t, attenMaskSize)
		40	+TILING_DATA_FIELD_DEF(uint32_t, sparseMode)
40	TILING_DATA_FIELD_DEF(uint32_t, outputLayout) // 输出Layout	41	TILING_DATA_FIELD_DEF(uint32_t, outputLayout) // 输出Layout
41	TILING_DATA_FIELD_DEF(uint32_t, softmaxLseFlag)	42	TILING_DATA_FIELD_DEF(uint32_t, softmaxLseFlag)
		43	+TILING_DATA_FIELD_DEF(uint32_t, tSeqSize) // TND T轴总长度（含padding）
42	END_TILING_DATA_DEF	44	END_TILING_DATA_DEF
43	REGISTER_TILING_DATA_CLASS(IncreFlashAttentionBaseParamsMlaOp, IncreFlashAttentionBaseParamsMla)	45	REGISTER_TILING_DATA_CLASS(IncreFlashAttentionBaseParamsMlaOp, IncreFlashAttentionBaseParamsMla)
44		46

Mattention/incre_flash_attention/op_kernel/arch32/incre_flash_attention_preload_mla.h+210-2

	@@ -336,6 +336,7 @@ protected:
336	uint64_t mSizeVStart = 0ULL;	336	uint64_t mSizeVStart = 0ULL;
337	uint64_t kvSeqSize = 0ULL;	337	uint64_t kvSeqSize = 0ULL;
338	uint64_t qSeqSize = 1ULL;	338	uint64_t qSeqSize = 1ULL;
		339	+ uint64_t tSeqSize = 1ULL; // TND T轴总长度（含padding）
339		340
340	// pageAttention	341	// pageAttention
341	uint32_t kvCacheBlockSize = 0;	342	uint32_t kvCacheBlockSize = 0;
	@@ -347,6 +348,7 @@ protected:
347	// attention mask	348	// attention mask
348	bool attenMaskFlag = false;	349	bool attenMaskFlag = false;
349	uint32_t attenMaskSizeAlign = 0U;	350	uint32_t attenMaskSizeAlign = 0U;
		351	+ uint32_t sparseMode = 3U;
350		352
351	// offset	353	// offset
352	uint64_t tensorACoreOffset = 0ULL;	354	uint64_t tensorACoreOffset = 0ULL;
	@@ -531,6 +533,9 @@ protected:
531	__aicore__ inline bool IsSkipAttenMask(const ExtraInfoMla &info, uint32_t startRow, uint32_t dealRowCount);	533	__aicore__ inline bool IsSkipAttenMask(const ExtraInfoMla &info, uint32_t startRow, uint32_t dealRowCount);
532	// 针对BSND/BSH/TND等切G的格式拷贝AttentionMask	534	// 针对BSND/BSH/TND等切G的格式拷贝AttentionMask
533	__aicore__ inline void AttenMaskCopyForSplitG(const ExtraInfoMla &info, LocalTensor<bool> &attenMaskUb, uint32_t startRow, uint32_t dealRowCount);	535	__aicore__ inline void AttenMaskCopyForSplitG(const ExtraInfoMla &info, LocalTensor<bool> &attenMaskUb, uint32_t startRow, uint32_t dealRowCount);
		536	+ __aicore__ inline void AttenMaskCopyForTree(const ExtraInfoMla &info, LocalTensor<bool> &attenMaskUb, uint32_t startRow, uint32_t dealRowCount);
		537	+ __aicore__ inline void AttenMaskCopyNoFull(LocalTensor<bool> &attenMaskUb, const ExtraInfoMla &info, uint32_t s1StartIdx, uint32_t s1EndIdx);
		538	+
534	__aicore__ inline void ElewiseCompute(const ExtraInfoMla &info, LocalTensor<T> &mmResUb, TBuf<> &tmpBuf, uint32_t startRow,	539	__aicore__ inline void ElewiseCompute(const ExtraInfoMla &info, LocalTensor<T> &mmResUb, TBuf<> &tmpBuf, uint32_t startRow,
535	uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount);	540	uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount);
536		541
	@@ -562,6 +567,7 @@ protected:
562		567
563	__aicore__ inline void InitAllZeroOutput(uint32_t bIdx, uint32_t n2Idx);	568	__aicore__ inline void InitAllZeroOutput(uint32_t bIdx, uint32_t n2Idx);
564	__aicore__ inline void InitSoftmaxLseAllInfOutput(uint32_t bIdx, uint32_t n2Idx);	569	__aicore__ inline void InitSoftmaxLseAllInfOutput(uint32_t bIdx, uint32_t n2Idx);
		570	+ __aicore__ inline void InitOutputSingleCore();
565	__aicore__ inline uint64_t SeqLenFromTensorList(uint32_t bIdx);	571	__aicore__ inline uint64_t SeqLenFromTensorList(uint32_t bIdx);
566		572
567	__aicore__ inline void CopyFixedUbToGm(const GlobalTensor<T> &dst, const LocalTensor<T> &src, size_t size);	573	__aicore__ inline void CopyFixedUbToGm(const GlobalTensor<T> &dst, const LocalTensor<T> &src, size_t size);
	@@ -590,6 +596,7 @@ template <typename IFAT> __aicore__ inline void IncreFlashAttentionAttenPreloadM
590		596
591	kvSeqSize = tilingData->baseParams.seqSize;	597	kvSeqSize = tilingData->baseParams.seqSize;
592	qSeqSize = tilingData->baseParams.qSeqSize;	598	qSeqSize = tilingData->baseParams.qSeqSize;
		599	+ tSeqSize = tilingData->baseParams.tSeqSize;
593		600
594	s1SizeSub = tilingData->increFlashAttentionSingleCoreParams.s1SplitSize; // 切块大小Si	601	s1SizeSub = tilingData->increFlashAttentionSingleCoreParams.s1SplitSize; // 切块大小Si
595	s1Outer = (qSeqSize + s1SizeSub - 1) / s1SizeSub;	602	s1Outer = (qSeqSize + s1SizeSub - 1) / s1SizeSub;
	@@ -597,6 +604,7 @@ template <typename IFAT> __aicore__ inline void IncreFlashAttentionAttenPreloadM
597		604
598	attenMaskFlag = (tilingData->baseParams.attenMaskFlag != 0) ? true : false;	605	attenMaskFlag = (tilingData->baseParams.attenMaskFlag != 0) ? true : false;
599	attenMaskSize = tilingData->baseParams.attenMaskSize;	606	attenMaskSize = tilingData->baseParams.attenMaskSize;
		607	+ sparseMode = tilingData->baseParams.sparseMode;
600		608
601	softmaxLseFlag = tilingData->baseParams.softmaxLseFlag;	609	softmaxLseFlag = tilingData->baseParams.softmaxLseFlag;
602		610
	@@ -732,6 +740,50 @@ __aicore__ inline void IncreFlashAttentionAttenPreloadMla<IFAT>::InitSoftmaxLseA
732	}	740	}
733	}	741	}
734		742
		743	+template <typename IFAT>
		744	+__aicore__ inline void IncreFlashAttentionAttenPreloadMla<IFAT>::InitOutputSingleCore()
		745	+{
		746	+ if (usedCoreNum != 0) {
		747	+ uint32_t initOutputEventId = 0U;
		748	+ SetFlag<AscendC::HardEvent::MTE3_V>(initOutputEventId);
		749	+ // TND 场景：sparse9 有 s2=0 的 batch，其 Q token 不会被计算路径写入，必须清零
		750	+ if constexpr (LAYOUT_T == LAYOUT::TND) {
		751	+ if (sparseMode == 9U && actualLenDims > 1U) {
		752	+ WaitFlag<AscendC::HardEvent::MTE3_V>(initOutputEventId);
		753	+ // 按 batch 轮询分配，跳过 s2≠0 的 batch
		754	+ for (uint64_t bIdx = (uint64_t)tmpBlockIdx; bIdx < batchSize;
		755	+ bIdx += (uint64_t)usedCoreNum) {
		756	+ uint64_t s2 = actualSeqLengthsGm.GetValue(bIdx);
		757	+ if (s2 != 0) {
		758	+ continue;
		759	+ }
		760	+ uint64_t tBase = (bIdx == 0) ? 0 : actualSeqLengthsGmQ.GetValue(bIdx - 1);
		761	+ uint64_t s1 = actualSeqLengthsGmQ.GetValue(bIdx) - tBase;
		762	+ if (s1 == 0) {
		763	+ continue;
		764	+ }
		765	+ // TND: [T, N, D]，batch 对应的 T 区间连续，一次清零
		766	+ uint64_t outputOffset = tBase * qHeadNum * headDim;
		767	+ uint64_t outputSize = s1 * qHeadNum * headDim;
		768	+ if constexpr (IsSameType<OUT_T, int8_t>::value) {
		769	+ GlobalTensor<half> attentionOutTmpGm;
		770	+ attentionOutTmpGm.SetGlobalBuffer(
		771	+ reinterpret_cast<__gm__ half *>(attentionOutGm.GetPhyAddr(0)));
		772	+ matmul::InitOutput<half>(
		773	+ attentionOutTmpGm[outputOffset / 2], outputSize / 2, 0);
		774	+ } else {
		775	+ matmul::InitOutput<OUT_T>(
		776	+ attentionOutGm[outputOffset], outputSize, 0);
		777	+ }
		778	+ }
		779	+ SetFlag<AscendC::HardEvent::MTE3_V>(initOutputEventId);
		780	+ }
		781	+ }
		782	+ WaitFlag<AscendC::HardEvent::MTE3_V>(initOutputEventId);
		783	+ SyncAll();
		784	+ }
		785	+}
		786	+
735	template <typename IFAT>	787	template <typename IFAT>
736	__aicore__ inline void IncreFlashAttentionAttenPreloadMla<IFAT>::GetActualSeqLen(uint32_t bIdx, uint32_t s1Idx)	788	__aicore__ inline void IncreFlashAttentionAttenPreloadMla<IFAT>::GetActualSeqLen(uint32_t bIdx, uint32_t s1Idx)
737	{	789	{
	@@ -895,6 +947,11 @@ __aicore__ inline void IncreFlashAttentionAttenPreloadMla<IFAT>::Init(
895		947
896	InitActualSeqLen(actualSeqLengthsQ, actualSeqLengths);	948	InitActualSeqLen(actualSeqLengthsQ, actualSeqLengths);
897		949
		950	+ // sparse9 TND padding 场景初始化 output 为 0
		951	+ if ASCEND_IS_AIV {
		952	+ InitOutputSingleCore();
		953	+ }
		954	+
898	if constexpr (PAGE_ATTENTION) {	955	if constexpr (PAGE_ATTENTION) {
899	blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable);	956	blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable);
900	}	957	}
	@@ -1880,6 +1937,144 @@ IncreFlashAttentionAttenPreloadMla<IFAT>::AttenMaskCopyForSplitG(const ExtraInfo
1880	attenMaskUb = attenMaskUbDst.template ReinterpretCast<bool>();	1937	attenMaskUb = attenMaskUbDst.template ReinterpretCast<bool>();
1881	}	1938	}
1882		1939
		1940	+// 新增sparse9的处理
		1941	+template <typename IFAT>
		1942	+__aicore__ inline void
		1943	+IncreFlashAttentionAttenPreloadMla<IFAT>::AttenMaskCopyNoFull(LocalTensor<bool> &attenMaskUb, const ExtraInfoMla &info, uint32_t s1StartIdx, uint32_t s1EndIdx)
		1944	+{
		1945	+ uint32_t actualSeqQ = qSeqSize;
		1946	+ if constexpr (LAYOUT_T == LAYOUT::TND) {
		1947	+ actualSeqQ = info.actS1Size;
		1948	+ }
		1949	+
		1950	+ uint32_t treeMaskStart = info.s2Size - actualSeqQ;
		1951	+ uint32_t curS2EndPos = info.s2Idx * singleProcessSInnerSize + info.actualSingleProcessSInnerSize;
		1952	+ uint32_t curS2StartPos = info.s2Idx * singleProcessSInnerSize;
		1953	+
		1954	+ uint64_t attenMaskBatchStride = actualSeqQ * actualSeqQ;
		1955	+ uint64_t attenMaskStride = actualSeqQ;
		1956	+
		1957	+ uint64_t bOffset;
		1958	+ uint64_t s1Offset;
		1959	+ if constexpr (LAYOUT_T == LAYOUT::TND) {
		1960	+ // TND格式下mask是一维张量∑s1²，batch偏移是前面所有batch的s1²累积和
		1961	+ uint64_t batchOffset = 0;
		1962	+ for (uint32_t i = 0; i < info.bIdx; i++) {
		1963	+ uint64_t s1i = GetBalanceActualSeqLengths(actualSeqLengthsGmQ, i);
		1964	+ batchOffset += s1i * s1i;
		1965	+ }
		1966	+ bOffset = batchOffset;
		1967	+ // TND格式下s1StartIdx是当前S1块内的局部偏移，完整局部S1索引需加上s1Idx*s1SizeSub
		1968	+ s1Offset = (static_cast<uint64_t>(info.s1Idx) * s1SizeSub + s1StartIdx) * attenMaskStride;
		1969	+ } else {
		1970	+ bOffset = static_cast<uint64_t>(info.bIdx) * static_cast<uint64_t>(attenMaskBatchStride);
		1971	+ s1Offset = (s1StartIdx % actualSeqQ) * attenMaskStride;
		1972	+ }
		1973	+ uint64_t s2Offset = curS2StartPos > treeMaskStart ? curS2StartPos - treeMaskStart : 0;
		1974	+
		1975	+ uint64_t maskOffset = bOffset + s1Offset + s2Offset;
		1976	+
		1977	+ if (curS2StartPos < treeMaskStart) { //mask占不满UB空间
		1978	+ // 不对齐场景下的搬运
		1979	+ DataCopyExtParams dataCopyParams;
		1980	+ dataCopyParams.blockCount = s1EndIdx - s1StartIdx;
		1981	+ dataCopyParams.blockLen = curS2EndPos - treeMaskStart;
		1982	+ dataCopyParams.srcStride = attenMaskStride - (curS2EndPos - treeMaskStart);
		1983	+ dataCopyParams.dstStride = (treeMaskStart - curS2StartPos) / 32;
		1984	+
		1985	+ DataCopyPadExtParams<bool> padParams;
		1986	+ padParams.isPad = true;
		1987	+ padParams.leftPadding = static_cast<uint8_t>(treeMaskStart % 32);
		1988	+ padParams.rightPadding = static_cast<uint8_t>(Align(static_cast<uint32_t>(curS2EndPos - treeMaskStart + treeMaskStart % 32), 32U) - \
		1989	+ (curS2EndPos - treeMaskStart + treeMaskStart % 32));
		1990	+ padParams.paddingValue = 0;
		1991	+ DataCopyPad(attenMaskUb[(treeMaskStart - curS2StartPos) / 32 * 32], attenMaskBoolGm[maskOffset], dataCopyParams, padParams);
		1992	+ } else { // mask 占满UB空间
		1993	+ DataCopyExtParams dataCopyParams;
		1994	+ dataCopyParams.blockCount = s1EndIdx - s1StartIdx;
		1995	+ dataCopyParams.blockLen = info.actualSingleProcessSInnerSize;
		1996	+ dataCopyParams.srcStride = attenMaskStride - info.actualSingleProcessSInnerSize;
		1997	+ dataCopyParams.dstStride = 0;
		1998	+ DataCopyPadExtParams<bool> padParams{true, 0, static_cast<uint8_t>(info.actualSingleProcessSInnerSizeAlign - info.actualSingleProcessSInnerSize), 0};
		1999	+ DataCopyPad(attenMaskUb, attenMaskBoolGm[maskOffset], dataCopyParams, padParams);
		2000	+ }
		2001	+}
		2002	+
		2003	+template <typename IFAT>
		2004	+__aicore__ inline void IncreFlashAttentionAttenPreloadMla<IFAT>::AttenMaskCopyForTree(const ExtraInfoMla &info,
		2005	+ LocalTensor<bool> &attenMaskUb, uint32_t startRow, uint32_t dealRowCount)
		2006	+{
		2007	+ uint32_t s1StartIdx = (mSizeVStart + startRow) / info.gSize;
		2008	+ uint32_t s1EndIdx = (mSizeVStart + startRow + dealRowCount - 1) / info.gSize;
		2009	+ uint32_t s1Count = s1EndIdx - s1StartIdx + 1;
		2010	+
		2011	+ uint32_t actualSeqQ = qSeqSize;
		2012	+ if constexpr (LAYOUT_T == LAYOUT::TND) {
		2013	+ actualSeqQ = info.actS1Size;
		2014	+ }
		2015	+
		2016	+
		2017	+ attenMaskSizeAlign = Align(info.actualSingleProcessSInnerSize, 32U);
		2018	+ // 第一步，将mask的 ub空间赋值为0
		2019	+ // 申请的16k空间，前8k给从GM拷贝到Ub时使用，后8k用来在拓展G轴时使用
		2020	+ attenMaskUb = inputQue2.AllocTensor<bool>();
		2021	+ LocalTensor<int16_t> mask16_0 = attenMaskUb.template ReinterpretCast<int16_t>();
		2022	+ AscendC::Duplicate(mask16_0, static_cast<int16_t>(0), s1Count * attenMaskSizeAlign / sizeof(int16_t));
		2023	+ attenMaskUb = mask16_0.template ReinterpretCast<bool>();
		2024	+ // V→MTE2 跨流水线同步：确保 V pipeline 的 Duplicate 完成后，MTE2 的 DataCopyPad 才能写入
		2025	+ event_t eventIdVMte2 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2));
		2026	+ SetFlag<HardEvent::V_MTE2>(eventIdVMte2);
		2027	+ WaitFlag<HardEvent::V_MTE2>(eventIdVMte2);
		2028	+
		2029	+ // 第二步，计算偏移从Gm上把mask拷过来
		2030	+ AttenMaskCopyNoFull(attenMaskUb, info, s1StartIdx, s1EndIdx + 1);
		2031	+
		2032	+ // 第三步，添加拷贝同步信号
		2033	+ inputQue2.template EnQue(attenMaskUb);
		2034	+ attenMaskUb = inputQue2.DeQue<bool>();
		2035	+
		2036	+ LocalTensor<int16_t> mask16 = attenMaskUb.template ReinterpretCast<int16_t>();
		2037	+ LocalTensor<int16_t> attenMaskUbDst = mask16[BUFFER_SIZE_BYTE_16K / 4];
		Y Y yang-binrong2月27日【提示】4等数组需要加注释，防止引入告警
		2038	+
		2039	+ uint32_t headGCount = 0;
		2040	+ uint32_t firstGIdx = (mSizeVStart + startRow) % info.gSize;
		2041	+ if (s1Count > 1) {
		2042	+ headGCount = info.gSize - firstGIdx;
		2043	+ } else {
		2044	+ headGCount = dealRowCount;
		2045	+ }
		2046	+ uint64_t dstMaskOffset = 0;
		2047	+ uint64_t srcMaskBaseOffset = 0;
		2048	+ // head
		2049	+ SetMaskCount();
		2050	+ SetVectorMask<int16_t, MaskMode::COUNTER>(attenMaskSizeAlign / 2);
		2051	+ Copy<int16_t, false>(attenMaskUbDst[dstMaskOffset], mask16[srcMaskBaseOffset],
		2052	+ AscendC::MASK_PLACEHOLDER, headGCount,
		2053	+ {1, 1, static_cast<uint16_t>(attenMaskSizeAlign / 32), 0});
		2054	+ dstMaskOffset += headGCount * attenMaskSizeAlign / 2;
		2055	+ srcMaskBaseOffset += attenMaskSizeAlign / 2;
		2056	+ // mid
		2057	+ uint32_t reminRowCount = dealRowCount - headGCount;
		2058	+ uint32_t midS1Count = reminRowCount / info.gSize;
		2059	+ uint32_t tailGSize = reminRowCount % info.gSize;
		2060	+ for (uint32_t midIdx = 0; midIdx < midS1Count; midIdx++) {
		2061	+ Copy<int16_t, false>(attenMaskUbDst[dstMaskOffset], mask16[srcMaskBaseOffset],
		2062	+ AscendC::MASK_PLACEHOLDER, info.gSize,
		2063	+ {1, 1, static_cast<uint16_t>(attenMaskSizeAlign / 32), 0});
		Y Y yang-binrong2月27日【提示】32等数字会引入告警，加注释
		2064	+ dstMaskOffset += info.gSize * attenMaskSizeAlign / 2;
		2065	+ srcMaskBaseOffset += attenMaskSizeAlign / 2;
		2066	+ }
		2067	+ // tail
		2068	+ if (tailGSize > 0) {
		2069	+ Copy<int16_t, false>(attenMaskUbDst[dstMaskOffset], mask16[srcMaskBaseOffset],
		2070	+ AscendC::MASK_PLACEHOLDER, tailGSize,
		2071	+ {1, 1, static_cast<uint16_t>(attenMaskSizeAlign / 32), 0});
		2072	+ }
		2073	+ SetMaskNorm();
		2074	+ ResetMask();
		2075	+ attenMaskUb = attenMaskUbDst.template ReinterpretCast<bool>();
		2076	+}
		2077	+
1883	template <typename IFAT>	2078	template <typename IFAT>
1884	__aicore__ inline bool	2079	__aicore__ inline bool
1885	IncreFlashAttentionAttenPreloadMla<IFAT>::IsSkipAttenMask(const ExtraInfoMla &info, uint32_t startRow, uint32_t dealRowCount)	2080	IncreFlashAttentionAttenPreloadMla<IFAT>::IsSkipAttenMask(const ExtraInfoMla &info, uint32_t startRow, uint32_t dealRowCount)
	@@ -1889,13 +2084,22 @@ IncreFlashAttentionAttenPreloadMla<IFAT>::IsSkipAttenMask(const ExtraInfoMla &in
1889	actualSeqQ = info.actS1Size;	2084	actualSeqQ = info.actS1Size;
1890	}	2085	}
1891		2086
		2087	+ uint32_t s2EndPos = info.s2Idx * singleProcessSInnerSize + info.actualSingleProcessSInnerSize;
		2088	+ // 新增sparse9的处理，sparse9和sparse3类似，但是sparse9是用户自己传入的mask矩阵，算子无法感知具体数值，所以涉及到基本块都要参与计算
		2089	+ if (sparseMode == 9U) {
		2090	+ if (s2EndPos > (info.s2Size - actualSeqQ)) {
		2091	+ return false;
		2092	+ } else {
		2093	+ return true;
		2094	+ }
		2095	+ }
		2096	+
1892	// s2<s1时，必然走mask	2097	// s2<s1时，必然走mask
1893	if (info.s2Size < actualSeqQ) {	2098	if (info.s2Size < actualSeqQ) {
1894	return false;	2099	return false;
1895	}	2100	}
1896		2101
1897	// 当前的s2位置不超过需要打标记的位置时，不需要mask	2102	// 当前的s2位置不超过需要打标记的位置时，不需要mask
1898	- uint32_t s2EndPos = info.s2Idx * singleProcessSInnerSize + info.actualSingleProcessSInnerSize;
1899	if (s2EndPos <= (info.s2Size - actualSeqQ + 1)) {	2103	if (s2EndPos <= (info.s2Size - actualSeqQ + 1)) {
1900	return true;	2104	return true;
1901	}	2105	}
	@@ -1966,7 +2170,11 @@ IncreFlashAttentionAttenPreloadMla<IFAT>::ElewiseCompute(const ExtraInfoMla &inf
1966	}	2170	}
1967	attenMaskUb = attenMaskUbDst;	2171	attenMaskUb = attenMaskUbDst;
1968	} else { // BSH/BSND/TND	2172	} else { // BSH/BSND/TND
1969	- AttenMaskCopyForSplitG(info, attenMaskUb, startRow, dealRowCount);	2173	+ if (sparseMode == 9U) {
		2174	+ AttenMaskCopyForTree(info, attenMaskUb, startRow, dealRowCount);
		2175	+ } else {
		2176	+ AttenMaskCopyForSplitG(info, attenMaskUb, startRow, dealRowCount);
		2177	+ }
1970	}	2178	}
1971	}	2179	}
1972		2180

Mdocs/zh/context/sparse_mode参数说明.md+26-1

	@@ -15,6 +15,7 @@
15	\| 6 \| prefix压缩模式。 \| - \|	15	\| 6 \| prefix压缩模式。 \| - \|
16	\| 7 \| varlen外切场景，rightDownCausal模式。 \| 仅varlen场景支持。 \|	16	\| 7 \| varlen外切场景，rightDownCausal模式。 \| 仅varlen场景支持。 \|
17	\| 8 \| varlen外切场景，leftUpCausal模式。 \| 仅varlen场景支持。 \|	17	\| 8 \| varlen外切场景，leftUpCausal模式。 \| 仅varlen场景支持。 \|
		18	+\| 9 \| treeMask模式。 \| 非量化支持GQA和MLA场景，全量化仅支持MLA场景。 \|
18		19
19	attenMask的工作原理为，在Mask为True的位置遮蔽query(Q)与key(K)的转置矩阵乘积的值，示意如下：	20	attenMask的工作原理为，在Mask为True的位置遮蔽query(Q)与key(K)的转置矩阵乘积的值，示意如下：
20		21
	@@ -146,4 +147,28 @@ Masked $QK^T$矩阵示意如下，在第二个batch对query进行切分，key和
146	- 基于sparseMode=2进行外切产生的band模式的sparse的参数应符合以下条件：	147	- 基于sparseMode=2进行外切产生的band模式的sparse的参数应符合以下条件：
147	- preTokens >= first_Skv。	148	- preTokens >= first_Skv。
148	- nextTokens >= first_Sq - first_Skv，根据实际情况进行配置。	149	- nextTokens >= first_Sq - first_Skv，根据实际情况进行配置。
149	- - 当前模式下不支持可选输入pse。	150	+ - 当前模式下不支持可选输入pse。
		151	+
		152	+## sparseMode=9
		153	+
		154	+sparseMode为9时，代表treeMask模式，用于推测解码（speculative decoding）场景下的树形注意力掩码。用户需传入自定义的树形mask，mask中值为1的位置会被遮蔽。
		155	+
		156	+树形mask矩阵特征如下：
		157	+- 对角线位置（s1==s2）：值为0，表示token关注自身。
		158	+- 上三角位置（s1<s2）：值为1，表示不关注未来token。
		159	+- 下三角位置（s1>s2）：值为0或1，由树结构决定部分注意力关系。
		160	+
		161	+attenMask输入格式：
		162	+- inputLayout为BSH、BSND或BNSD时：attenMask的shape为(B, S1, S1)，每个batch传入S1×S1大小的tree mask。
		163	+
		164	+ ![原理图](../figures/sparsemode为9_BSND_BNSD示意图.png)
		165	+
		166	+- inputLayout为TND时：attenMask为1D紧凑格式，shape为(∑S1i²,)，即每个batch的S1i×S1i mask拼接传入。
		167	+
		168	+ ![原理图](../figures/sparsemode为9_TND示意图.png)
		169	+
		170	+约束说明：
		171	+- 非量化场景支持GQA和MLA，全量化场景仅支持MLA。
		172	+- 不支持左padding、pseShift、sharedPrefix。
		173	+- 输出dtype不支持INT8。
		174	+- 每个batch需满足Q_S ≤ KV_S。

Adocs/zh/figures/sparsemode为9_BSND_BNSD示意图.png+0-0

Adocs/zh/figures/sparsemode为9_TND示意图.png+0-0