同步更新19MLA的optest测试件-catlass-AtomGit

已合并

同步更新19MLA的optest测试件 #1163

init__zhb__创建于 22 天前从init__zhb__/catlass_zhb： br.19mla_update 合入到 cann/catlass： master

同步更新19MLA的optest测试件 #1163

已合并

init__zhb__创建于 22 天前

从init__zhb__/catlass_zhb br.19mla_update合入到cann/catlass master

共 2 个文件变更+222-144

Mtests/optest/kernels/19_mla/mla_kernel.cpp+146-135

	@@ -33,10 +33,7 @@ using namespace Catlass;
33	This example demonstrates how to compute mla.	33	This example demonstrates how to compute mla.
34	*/	34	*/
35	template <	35	template <
36	- class BlockMmadQK,	36	+ class BlockMmadQK, class BlockMmadPV, class EpilogueMLASoftmax, class EpilogueMLARescaleO,
37	- class BlockMmadPV,
38	- class EpilogueMLASoftmax,
39	- class EpilogueMLARescaleO,
40	class EpilogueMLAFDRescaleO>	37	class EpilogueMLAFDRescaleO>
41	class MLAKernel {	38	class MLAKernel {
42	public:	39	public:
	@@ -88,21 +85,35 @@ public:
88		85
89	// Methods	86	// Methods
90	CATLASS_DEVICE	87	CATLASS_DEVICE
91	- Params() {}	88	+ Params()
		89	+ {}
92		90
93	CATLASS_DEVICE	91	CATLASS_DEVICE
94	- Params(GM_ADDR q_, GM_ADDR qRope_, GM_ADDR k_, GM_ADDR kRope_, GM_ADDR blockTables_,	92	+ Params(
95	- GM_ADDR o_, GM_ADDR s_, GM_ADDR p_, GM_ADDR oTmp_, GM_ADDR oUpdate_,	93	+ GM_ADDR q_, GM_ADDR qRope_, GM_ADDR k_, GM_ADDR kRope_, GM_ADDR blockTables_, GM_ADDR o_, GM_ADDR s_,
96	- GM_ADDR oCoreTmp_, GM_ADDR l_, GM_ADDR tiling_)	94	+ GM_ADDR p_, GM_ADDR oTmp_, GM_ADDR oUpdate_, GM_ADDR oCoreTmp_, GM_ADDR l_, GM_ADDR tiling_)
97	- : q(q_), qRope(qRope_), k(k_), kRope(kRope_), blockTables(blockTables_), o(o_),	95	+ : q(q_),
98	- s(s_), p(p_), oTmp(oTmp_), oUpdate(oUpdate_), oCoreTmp(oCoreTmp_), l(l_), tiling(tiling_) {}	96	+ qRope(qRope_),
		97	+ k(k_),
		98	+ kRope(kRope_),
		99	+ blockTables(blockTables_),
		100	+ o(o_),
		101	+ s(s_),
		102	+ p(p_),
		103	+ oTmp(oTmp_),
		104	+ oUpdate(oUpdate_),
		105	+ oCoreTmp(oCoreTmp_),
		106	+ l(l_),
		107	+ tiling(tiling_)
		108	+ {}
99	};	109	};
100		110
101	// Methods	111	// Methods
102	CATLASS_DEVICE	112	CATLASS_DEVICE
103	- MLAKernel() {}	113	+ MLAKernel()
		114	+ {}
104		115
105	- CATLASS_DEVICE void operator()(Params const &params)	116	+ CATLASS_DEVICE void operator()(Params const& params)
106	{	117	{
107	#ifdef __DAV_CUBE__	118	#ifdef __DAV_CUBE__
108	AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);	119	AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
	@@ -151,41 +162,41 @@ public:
151	#endif	162	#endif
152		163
153	AscendC::GlobalTensor<ElementS> gS;	164	AscendC::GlobalTensor<ElementS> gS;
154	- gS.SetGlobalBuffer((__gm__ ElementS *)params.s);	165	+ gS.SetGlobalBuffer((__gm__ ElementS*)params.s);
155	AscendC::GlobalTensor<ElementP> gP;	166	AscendC::GlobalTensor<ElementP> gP;
156	- gP.SetGlobalBuffer((__gm__ ElementP *)params.p);	167	+ gP.SetGlobalBuffer((__gm__ ElementP*)params.p);
157	AscendC::GlobalTensor<ElementOTmp> gOTmp;	168	AscendC::GlobalTensor<ElementOTmp> gOTmp;
158	- gOTmp.SetGlobalBuffer((__gm__ ElementOTmp *)params.oTmp);	169	+ gOTmp.SetGlobalBuffer((__gm__ ElementOTmp*)params.oTmp);
159	AscendC::GlobalTensor<uint32_t> gTiling;	170	AscendC::GlobalTensor<uint32_t> gTiling;
160	- gTiling.SetGlobalBuffer((__gm__ uint32_t *)params.tiling);	171	+ gTiling.SetGlobalBuffer((__gm__ uint32_t*)params.tiling);
161		172
162	#ifdef __DAV_CUBE__	173	#ifdef __DAV_CUBE__
163	// Get the memory offset address of the input on Global Memory	174	// Get the memory offset address of the input on Global Memory
164	AscendC::GlobalTensor<ElementQ> gQ;	175	AscendC::GlobalTensor<ElementQ> gQ;
165	- gQ.SetGlobalBuffer((__gm__ ElementQ *)params.q);	176	+ gQ.SetGlobalBuffer((__gm__ ElementQ*)params.q);
166	AscendC::GlobalTensor<ElementQ> gQRope;	177	AscendC::GlobalTensor<ElementQ> gQRope;
167	- gQRope.SetGlobalBuffer((__gm__ ElementQ *)params.qRope);	178	+ gQRope.SetGlobalBuffer((__gm__ ElementQ*)params.qRope);
168	AscendC::GlobalTensor<ElementK> gK;	179	AscendC::GlobalTensor<ElementK> gK;
169	- gK.SetGlobalBuffer((__gm__ ElementK *)params.k);	180	+ gK.SetGlobalBuffer((__gm__ ElementK*)params.k);
170	AscendC::GlobalTensor<ElementK> gKRope;	181	AscendC::GlobalTensor<ElementK> gKRope;
171	- gKRope.SetGlobalBuffer((__gm__ ElementK *)params.kRope);	182	+ gKRope.SetGlobalBuffer((__gm__ ElementK*)params.kRope);
172	AscendC::GlobalTensor<int32_t> gblockTable;	183	AscendC::GlobalTensor<int32_t> gblockTable;
173	- gblockTable.SetGlobalBuffer((__gm__ int32_t *)(params.blockTables));	184	+ gblockTable.SetGlobalBuffer((__gm__ int32_t*)(params.blockTables));
174		185
175	BlockMmadQK blockMmadQK(resource);	186	BlockMmadQK blockMmadQK(resource);
176	BlockMmadPV blockMmadPV(resource);	187	BlockMmadPV blockMmadPV(resource);
177	#endif	188	#endif
178	#ifdef __DAV_VEC__	189	#ifdef __DAV_VEC__
179	AscendC::GlobalTensor<ElementO> gO;	190	AscendC::GlobalTensor<ElementO> gO;
180	- gO.SetGlobalBuffer((__gm__ ElementO *)params.o);	191	+ gO.SetGlobalBuffer((__gm__ ElementO*)params.o);
181	AscendC::GlobalTensor<ElementOTmp> gOUpdate;	192	AscendC::GlobalTensor<ElementOTmp> gOUpdate;
182	- gOUpdate.SetGlobalBuffer((__gm__ ElementOTmp *)params.oUpdate);	193	+ gOUpdate.SetGlobalBuffer((__gm__ ElementOTmp*)params.oUpdate);
183	AscendC::GlobalTensor<ElementOTmp> gOCoreTmp;	194	AscendC::GlobalTensor<ElementOTmp> gOCoreTmp;
184	- gOCoreTmp.SetGlobalBuffer((__gm__ ElementOTmp *)params.oCoreTmp);	195	+ gOCoreTmp.SetGlobalBuffer((__gm__ ElementOTmp*)params.oCoreTmp);
185	AscendC::GlobalTensor<ElementOTmp> gl;	196	AscendC::GlobalTensor<ElementOTmp> gl;
186	- gl.SetGlobalBuffer((__gm__ ElementOTmp *)params.l);	197	+ gl.SetGlobalBuffer((__gm__ ElementOTmp*)params.l);
187	AscendC::GlobalTensor<float> gTilingFp64;	198	AscendC::GlobalTensor<float> gTilingFp64;
188	- gTilingFp64.SetGlobalBuffer((__gm__ float *)params.tiling);	199	+ gTilingFp64.SetGlobalBuffer((__gm__ float*)params.tiling);
189	#endif	200	#endif
190		201
191	// Get tiling parameters	202	// Get tiling parameters
	@@ -203,7 +214,7 @@ public:
203	uint32_t embedRound = RoundUp<BLOCK_SIZE>(embed);	214	uint32_t embedRound = RoundUp<BLOCK_SIZE>(embed);
204	#ifdef __DAV_VEC__	215	#ifdef __DAV_VEC__
205	float tor = gTilingFp64.GetValue(TILING_TOR);	216	float tor = gTilingFp64.GetValue(TILING_TOR);
206	- uint32_t glFlag[2] = {1,1};	217	+ uint32_t glFlag[2] = {1, 1};
207		218
208	EpilogueMLASoftmax epilogueMLASoftmax(resource, tor, maxKvSplitCoreNum);	219	EpilogueMLASoftmax epilogueMLASoftmax(resource, tor, maxKvSplitCoreNum);
209	EpilogueMLARescaleO epilogueMLARescaleO(resource, maxKvSplitCoreNum);	220	EpilogueMLARescaleO epilogueMLARescaleO(resource, maxKvSplitCoreNum);
	@@ -226,12 +237,18 @@ public:
226	coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum();	237	coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum();
227	#endif	238	#endif
228	uint32_t coreNum = AscendC::GetBlockNum();	239	uint32_t coreNum = AscendC::GetBlockNum();
229	- uint32_t processNum = batch * curQheadSplitNum * maxKvSplitCoreNum;	240	+ // When the tiling decides to split kv it publishes a non-zero TILING_PROCESSNUM together
		241	+ // with a cumulative-task prefix-sum table (CUTASK_START_OFFSET ...). In that case the
		242	+ // process -> task mapping is the dense T+1-style mapping; otherwise (no kv split) fall
		243	+ // back to the original back-and-forth (kerFlag / isForward) dispatch.
		244	+ uint32_t tilingProcessNum = gTiling.GetValue(TILING_PROCESSNUM);
		245	+ uint32_t processNum = (tilingProcessNum != 0) ? tilingProcessNum : batch * curQheadSplitNum * maxKvSplitCoreNum;
230		246
231	bool kerFlag = true;	247	bool kerFlag = true;
232	bool isForward = true;	248	bool isForward = true;
233	uint32_t pingpongIdx = 0;	249	uint32_t pingpongIdx = 0;
234	bool isFirstTask = true;	250	bool isFirstTask = true;
		251	+ uint32_t taskIdx = 0;
235	#ifdef __DAV_CUBE__	252	#ifdef __DAV_CUBE__
236	uint32_t locPingPongIdx = 0;	253	uint32_t locPingPongIdx = 0;
237	#endif	254	#endif
	@@ -241,12 +258,20 @@ public:
241		258
242	// Cube and vector share the same task traversal; stage bodies stay in macro branches.	259	// Cube and vector share the same task traversal; stage bodies stay in macro branches.
243	for (uint32_t process = coreIdx; process < processNum; process += uint32_t(coreNum)) {	260	for (uint32_t process = coreIdx; process < processNum; process += uint32_t(coreNum)) {
244	- uint32_t bigProcess = process - (process % coreNum) + (coreNum - 1);	261	+ uint32_t curBatch;
245	- bigProcess = (bigProcess > processNum - 1) ? (processNum - 1) : bigProcess;	262	+ if (tilingProcessNum != 0) {
246	- uint32_t realProcess = isForward ? process : (bigProcess - process % coreNum);	263	+ // T+1-style continuous (prefix-sum) task mapping
247	- isForward = !isForward;	264	+ while (taskIdx < batch && process >= gTiling.GetValue(CUTASK_START_OFFSET + taskIdx + 1)) {
248	-	265	+ taskIdx++;
249	- uint32_t curBatch = realProcess / (curQheadSplitNum * maxKvSplitCoreNum);	266	+ }
		267	+ curBatch = taskIdx;
		268	+ } else {
		269	+ uint32_t bigProcess = process - (process % coreNum) + (coreNum - 1);
		270	+ bigProcess = (bigProcess > processNum - 1) ? (processNum - 1) : bigProcess;
		271	+ uint32_t realProcess = isForward ? process : (bigProcess - process % coreNum);
		272	+ isForward = !isForward;
		273	+ curBatch = realProcess / (curQheadSplitNum * maxKvSplitCoreNum);
		274	+ }
250	uint32_t offsetTiling = tilingHeadSize + tilingParaSize * curBatch;	275	uint32_t offsetTiling = tilingHeadSize + tilingParaSize * curBatch;
251	uint32_t qSeqlen = gTiling.GetValue(offsetTiling);	276	uint32_t qSeqlen = gTiling.GetValue(offsetTiling);
252	uint32_t kvSeqlen = gTiling.GetValue(offsetTiling + 1);	277	uint32_t kvSeqlen = gTiling.GetValue(offsetTiling + 1);
	@@ -257,20 +282,26 @@ public:
257	continue;	282	continue;
258	}	283	}
259		284
260	- uint32_t qHeadSplitIdx = (process % (curQheadSplitNum * maxKvSplitCoreNum)) / maxKvSplitCoreNum;	285	+ uint32_t qHeadSplitIdx;
261	- uint32_t qHeadSplitSizeActual = (qHeadSplitIdx == (curQheadSplitNum - 1))	286	+ uint32_t curNIdx;
262	- ? (qHeads - qHeadSplitIdx * curQheadSplitSize)	287	+ if (tilingProcessNum != 0) {
263	- : curQheadSplitSize;	288	+ qHeadSplitIdx = 0;
264	- uint32_t curStartHeadIdx = qHeadSplitIdx * curQheadSplitSize;	289	+ curNIdx = process - gTiling.GetValue(CUTASK_START_OFFSET + curBatch);
265	- uint32_t curNIdx = process % maxKvSplitCoreNum;
266	- uint32_t curKVSeqlen = kvSplitPerCore;
267	-
268	- if (kerFlag) {
269	- kerFlag = false;
270	} else {	290	} else {
271	- kerFlag = true;	291	+ qHeadSplitIdx = (process % (curQheadSplitNum * maxKvSplitCoreNum)) / maxKvSplitCoreNum;
272	- curNIdx = maxKvSplitCoreNum - curNIdx - 1;	292	+ curNIdx = process % maxKvSplitCoreNum;
		293	+ if (kerFlag) {
		294	+ kerFlag = false;
		295	+ } else {
		296	+ kerFlag = true;
		297	+ curNIdx = maxKvSplitCoreNum - curNIdx - 1;
		298	+ }
273	}	299	}
		300	+ uint32_t qHeadSplitSizeActual = (qHeadSplitIdx == (curQheadSplitNum - 1)) ?
		301	+ (qHeads - qHeadSplitIdx * curQheadSplitSize) :
		302	+ curQheadSplitSize;
		303	+ uint32_t curStartHeadIdx = qHeadSplitIdx * curQheadSplitSize;
		304	+ uint32_t curKVSeqlen = kvSplitPerCore;
274		305
275	if (curNIdx >= kvSplitCoreNum) {	306	if (curNIdx >= kvSplitCoreNum) {
276	continue;	307	continue;
	@@ -344,14 +375,9 @@ public:
344	uint64_t gSOffset =	375	uint64_t gSOffset =
345	(uint64_t)coreIdx * TMP_SIZE_DECODER + (uint64_t)qkPingPongFlag * TMP_SIZE_DECODER / 2;	376	(uint64_t)coreIdx * TMP_SIZE_DECODER + (uint64_t)qkPingPongFlag * TMP_SIZE_DECODER / 2;
346	blockMmadQK(	377	blockMmadQK(
347	- gQ[gQOffset],	378	+ gQ[gQOffset], gQRope[gQRopeOffset], gK[kvOffset], gKRope[kvOffsetRope], gS[gSOffset], layoutQ,
348	- gQRope[gQRopeOffset],	379	+ layoutQRope, layoutK, layoutKRope, layoutS, actualBlockShapeQK, qShapeSingleNd, qHeads, nIdx,
349	- gK[kvOffset],	380	+ pingpongIdx);
350	- gKRope[kvOffsetRope],
351	- gS[gSOffset],
352	- layoutQ, layoutQRope, layoutK, layoutKRope, layoutS,
353	- actualBlockShapeQK, qShapeSingleNd,
354	- qHeads, nIdx, pingpongIdx);
355		381
356	pingpongIdx++;	382	pingpongIdx++;
357	Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady);	383	Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady);
	@@ -365,13 +391,10 @@ public:
365	LayoutS layoutS(rowNumRound, kSeqTile, kSeqTileRound);	391	LayoutS layoutS(rowNumRound, kSeqTile, kSeqTileRound);
366	GemmCoord actualBlockShapeQK{rowNum, kSeqTile, embedRound};	392	GemmCoord actualBlockShapeQK{rowNum, kSeqTile, embedRound};
367	uint64_t gmOffsetP = (uint64_t)coreIdx * TMP_SIZE + softmaxPingPongFlag * TMP_SIZE / 2;	393	uint64_t gmOffsetP = (uint64_t)coreIdx * TMP_SIZE + softmaxPingPongFlag * TMP_SIZE / 2;
368	- uint64_t gmOffsetS =	394	+ uint64_t gmOffsetS = (uint64_t)coreIdx * TMP_SIZE_DECODER + softmaxPingPongFlag * TMP_SIZE_DECODER / 2;
369	- (uint64_t)coreIdx * TMP_SIZE_DECODER + softmaxPingPongFlag * TMP_SIZE_DECODER / 2;
370	epilogueMLASoftmax(	395	epilogueMLASoftmax(
371	- gP[gmOffsetP], gS[gmOffsetS],	396	+ gP[gmOffsetP], gS[gmOffsetS], layoutP, layoutS, actualBlockShapeQK, nIdx, qHeadSplitSizeActual,
372	- layoutP, layoutS,	397	+ softmaxPingPongFlag, glFlag, taskPingPongFlag);
373	- actualBlockShapeQK,
374	- nIdx, qHeadSplitSizeActual, softmaxPingPongFlag, glFlag, taskPingPongFlag);
375		398
376	pingpongIdx++;	399	pingpongIdx++;
377	Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxReady);	400	Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxReady);
	@@ -401,14 +424,9 @@ public:
401	uint64_t gSOffset =	424	uint64_t gSOffset =
402	(uint64_t)coreIdx * TMP_SIZE_DECODER + (uint64_t)qkPingPongFlag * TMP_SIZE_DECODER / 2;	425	(uint64_t)coreIdx * TMP_SIZE_DECODER + (uint64_t)qkPingPongFlag * TMP_SIZE_DECODER / 2;
403	blockMmadQK(	426	blockMmadQK(
404	- gQ[gQOffset],	427	+ gQ[gQOffset], gQRope[gQRopeOffset], gK[kvOffset], gKRope[kvOffsetRope], gS[gSOffset], layoutQ,
405	- gQRope[gQRopeOffset],	428	+ layoutQRope, layoutK, layoutKRope, layoutS, actualBlockShapeQK, qShapeSingleNd, qHeads, nIdx,
406	- gK[kvOffset],	429	+ pingpongIdx);
407	- gKRope[kvOffsetRope],
408	- gS[gSOffset],
409	- layoutQ, layoutQRope, layoutK, layoutKRope, layoutS,
410	- actualBlockShapeQK, qShapeSingleNd,
411	- qHeads, nIdx, pingpongIdx);
412		430
413	Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady);	431	Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady);
414	#endif	432	#endif
	@@ -423,10 +441,8 @@ public:
423	uint64_t gmOffsetS =	441	uint64_t gmOffsetS =
424	(uint64_t)coreIdx * TMP_SIZE_DECODER + softmaxPingPongFlag * TMP_SIZE_DECODER / 2;	442	(uint64_t)coreIdx * TMP_SIZE_DECODER + softmaxPingPongFlag * TMP_SIZE_DECODER / 2;
425	epilogueMLASoftmax(	443	epilogueMLASoftmax(
426	- gP[gmOffsetP], gS[gmOffsetS],	444	+ gP[gmOffsetP], gS[gmOffsetS], layoutP, layoutS, actualBlockShapeQK, nIdx, qHeadSplitSizeActual,
427	- layoutP, layoutS,	445	+ softmaxPingPongFlag, glFlag, taskPingPongFlag);
428	- actualBlockShapeQK,
429	- nIdx, qHeadSplitSizeActual, softmaxPingPongFlag, glFlag, taskPingPongFlag);
430		446
431	Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxReady);	447	Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxReady);
432	#endif	448	#endif
	@@ -435,14 +451,24 @@ public:
435	if (nIdx == nLoop) {	451	if (nIdx == nLoop) {
436	bool nextKerFlag = kerFlag;	452	bool nextKerFlag = kerFlag;
437	bool nextIsForward = isForward;	453	bool nextIsForward = isForward;
		454	+ uint32_t nextTaskIdx = taskIdx;
438	for (uint32_t nextProcess = uint32_t(coreNum) + process; nextProcess < processNum;	455	for (uint32_t nextProcess = uint32_t(coreNum) + process; nextProcess < processNum;
439	nextProcess += uint32_t(coreNum)) {	456	nextProcess += uint32_t(coreNum)) {
440	- uint32_t nextBigProcess = nextProcess - (nextProcess % coreNum) + (coreNum - 1);	457	+ uint32_t nextBatch;
441	- nextBigProcess = (nextBigProcess > processNum - 1) ? (processNum - 1) : nextBigProcess;	458	+ if (tilingProcessNum != 0) {
442	- uint32_t nextRealProcess = nextIsForward ? nextProcess : (nextBigProcess - nextProcess % coreNum);	459	+ while (nextTaskIdx < batch &&
443	- nextIsForward = !nextIsForward;	460	+ nextProcess >= gTiling.GetValue(CUTASK_START_OFFSET + nextTaskIdx + 1)) {
444	-	461	+ nextTaskIdx++;
445	- uint32_t nextBatch = nextRealProcess / (curQheadSplitNum * maxKvSplitCoreNum);	462	+ }
		463	+ nextBatch = nextTaskIdx;
		464	+ } else {
		465	+ uint32_t nextBigProcess = nextProcess - (nextProcess % coreNum) + (coreNum - 1);
		466	+ nextBigProcess = (nextBigProcess > processNum - 1) ? (processNum - 1) : nextBigProcess;
		467	+ uint32_t nextRealProcess =
		468	+ nextIsForward ? nextProcess : (nextBigProcess - nextProcess % coreNum);
		469	+ nextIsForward = !nextIsForward;
		470	+ nextBatch = nextRealProcess / (curQheadSplitNum * maxKvSplitCoreNum);
		471	+ }
446	uint32_t nextOffsetTiling = tilingHeadSize + tilingParaSize * nextBatch;	472	uint32_t nextOffsetTiling = tilingHeadSize + tilingParaSize * nextBatch;
447	uint32_t nextQSeqlen = gTiling.GetValue(nextOffsetTiling);	473	uint32_t nextQSeqlen = gTiling.GetValue(nextOffsetTiling);
448	uint32_t nextKVSeqlen = gTiling.GetValue(nextOffsetTiling + 1);	474	uint32_t nextKVSeqlen = gTiling.GetValue(nextOffsetTiling + 1);
	@@ -453,20 +479,26 @@ public:
453	continue;	479	continue;
454	}	480	}
455		481
456	- uint32_t nextQHeadSplitIdx =	482	+ uint32_t nextQHeadSplitIdx;
457	- (nextProcess % (curQheadSplitNum * maxKvSplitCoreNum)) / maxKvSplitCoreNum;	483	+ uint32_t nextCurNIdx;
458	- uint32_t nextQHeadSplitSizeActual = (nextQHeadSplitIdx == (curQheadSplitNum - 1))	484	+ if (tilingProcessNum != 0) {
459	- ? (qHeads - nextQHeadSplitIdx * curQheadSplitSize)	485	+ nextQHeadSplitIdx = 0;
460	- : curQheadSplitSize;	486	+ nextCurNIdx = nextProcess - gTiling.GetValue(CUTASK_START_OFFSET + nextBatch);
461	- uint32_t nextCurNIdx = nextProcess % maxKvSplitCoreNum;
462	- uint32_t nextCurKVSeqlen = nextKVSplitPerCore;
463	-
464	- if (nextKerFlag) {
465	- nextKerFlag = false;
466	} else {	487	} else {
467	- nextKerFlag = true;	488	+ nextQHeadSplitIdx =
468	- nextCurNIdx = maxKvSplitCoreNum - nextCurNIdx - 1;	489	+ (nextProcess % (curQheadSplitNum * maxKvSplitCoreNum)) / maxKvSplitCoreNum;
		490	+ nextCurNIdx = nextProcess % maxKvSplitCoreNum;
		491	+ if (nextKerFlag) {
		492	+ nextKerFlag = false;
		493	+ } else {
		494	+ nextKerFlag = true;
		495	+ nextCurNIdx = maxKvSplitCoreNum - nextCurNIdx - 1;
		496	+ }
469	}	497	}
		498	+ uint32_t nextQHeadSplitSizeActual = (nextQHeadSplitIdx == (curQheadSplitNum - 1)) ?
		499	+ (qHeads - nextQHeadSplitIdx * curQheadSplitSize) :
		500	+ curQheadSplitSize;
		501	+ uint32_t nextCurKVSeqlen = nextKVSplitPerCore;
470		502
471	if (nextCurNIdx >= nextKVSplitCoreNum) {	503	if (nextCurNIdx >= nextKVSplitCoreNum) {
472	continue;	504	continue;
	@@ -508,21 +540,16 @@ public:
508	GemmCoord nextActualBlockShapeQK{nextRowNum, nextKSeqTile, embed + embedRope};	540	GemmCoord nextActualBlockShapeQK{nextRowNum, nextKSeqTile, embed + embedRope};
509	MatrixCoord nextQShapeSingleNd{nextQHeadSplitSizeActual, embed};	541	MatrixCoord nextQShapeSingleNd{nextQHeadSplitSizeActual, embed};
510	uint32_t qkPingPongFlag = pingpongIdx % 2;	542	uint32_t qkPingPongFlag = pingpongIdx % 2;
511	- int32_t nextBlockTableId =	543	+ int32_t nextBlockTableId = gblockTable.GetValue(
512	- gblockTable.GetValue(nextRealBatch * maxNumBlocksPerQuery + nextStartKV / blockSize + nextNIdx);	544	+ nextRealBatch * maxNumBlocksPerQuery + nextStartKV / blockSize + nextNIdx);
513	uint64_t nextKvOffset = (uint64_t)nextBlockTableId * blockSize * strideKV;	545	uint64_t nextKvOffset = (uint64_t)nextBlockTableId * blockSize * strideKV;
514	uint64_t nextKvOffsetRope = (uint64_t)nextBlockTableId * blockSize * strideKVRope;	546	uint64_t nextKvOffsetRope = (uint64_t)nextBlockTableId * blockSize * strideKVRope;
515	uint64_t nextGSOffset =	547	uint64_t nextGSOffset =
516	(uint64_t)coreIdx * TMP_SIZE_DECODER + (uint64_t)qkPingPongFlag * TMP_SIZE_DECODER / 2;	548	(uint64_t)coreIdx * TMP_SIZE_DECODER + (uint64_t)qkPingPongFlag * TMP_SIZE_DECODER / 2;
517	blockMmadQK(	549	blockMmadQK(
518	- gQ[nextGQOffset],	550	+ gQ[nextGQOffset], gQRope[nextGQRopeOffset], gK[nextKvOffset], gKRope[nextKvOffsetRope],
519	- gQRope[nextGQRopeOffset],	551	+ gS[nextGSOffset], nextLayoutQ, nextLayoutQRope, nextLayoutK, nextLayoutKRope, nextLayoutS,
520	- gK[nextKvOffset],	552	+ nextActualBlockShapeQK, nextQShapeSingleNd, qHeads, nextNIdx, pingpongIdx);
521	- gKRope[nextKvOffsetRope],
522	- gS[nextGSOffset],
523	- nextLayoutQ, nextLayoutQRope, nextLayoutK, nextLayoutKRope, nextLayoutS,
524	- nextActualBlockShapeQK, nextQShapeSingleNd,
525	- qHeads, nextNIdx, pingpongIdx);
526		553
527	Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady);	554	Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady);
528	#endif	555	#endif
	@@ -538,9 +565,7 @@ public:
538	uint64_t nextGMOffsetS =	565	uint64_t nextGMOffsetS =
539	(uint64_t)coreIdx * TMP_SIZE_DECODER + softmaxPingPongFlag * TMP_SIZE_DECODER / 2;	566	(uint64_t)coreIdx * TMP_SIZE_DECODER + softmaxPingPongFlag * TMP_SIZE_DECODER / 2;
540	epilogueMLASoftmax(	567	epilogueMLASoftmax(
541	- gP[nextGMOffsetP], gS[nextGMOffsetS],	568	+ gP[nextGMOffsetP], gS[nextGMOffsetS], nextLayoutP, nextLayoutS, nextActualBlockShapeQK,
542	- nextLayoutP, nextLayoutS,
543	- nextActualBlockShapeQK,
544	nextNIdx, nextQHeadSplitSizeActual, softmaxPingPongFlag, glFlag, nextTaskPingPongFlag);	569	nextNIdx, nextQHeadSplitSizeActual, softmaxPingPongFlag, glFlag, nextTaskPingPongFlag);
545		570
546	Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxReady);	571	Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxReady);
	@@ -563,10 +588,8 @@ public:
563	uint64_t gPOffset = (uint64_t)coreIdx * TMP_SIZE + (uint64_t)pvPingPongFlag * TMP_SIZE / 2;	588	uint64_t gPOffset = (uint64_t)coreIdx * TMP_SIZE + (uint64_t)pvPingPongFlag * TMP_SIZE / 2;
564	uint64_t gOTmpOffset = (uint64_t)coreIdx * TMP_SIZE * 2 + (uint64_t)pvPingPongFlag * TMP_SIZE;	589	uint64_t gOTmpOffset = (uint64_t)coreIdx * TMP_SIZE * 2 + (uint64_t)pvPingPongFlag * TMP_SIZE;
565	blockMmadPV(	590	blockMmadPV(
566	- gP[gPOffset],	591	+ gP[gPOffset], gOTmp[gOTmpOffset], layoutP, layoutV, layoutOTmp, actualBlockShapePV, nIdx,
567	- gOTmp[gOTmpOffset],	592	+ pingpongIdx, softmaxReady, locPingPongIdx);
568	- layoutP, layoutV, layoutOTmp,
569	- actualBlockShapePV, nIdx, pingpongIdx, softmaxReady, locPingPongIdx);
570	Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(pvReady);	593	Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(pvReady);
571	#endif	594	#endif
572	#ifdef __DAV_VEC__	595	#ifdef __DAV_VEC__
	@@ -581,11 +604,9 @@ public:
581	uint64_t gmOffsetUpdate = (uint64_t)(coreIdx * TMP_SIZE);	604	uint64_t gmOffsetUpdate = (uint64_t)(coreIdx * TMP_SIZE);
582	uint32_t isLastNTile = (nIdx == nLoop) ? 1 : 0;	605	uint32_t isLastNTile = (nIdx == nLoop) ? 1 : 0;
583	epilogueMLARescaleO(	606	epilogueMLARescaleO(
584	- gOTmp[gmOffsetOTmp], gOUpdate[gmOffsetUpdate], gO[gmOffsetO],	607	+ gOTmp[gmOffsetOTmp], gOUpdate[gmOffsetUpdate], gO[gmOffsetO], gOCoreTmp[oFdOffset], gl[lOffset],
585	- gOCoreTmp[oFdOffset], gl[lOffset],	608	+ layoutOTmp, layoutO, layoutUpdate, actualBlockShapePV, nIdx, isLastNTile, qHeadSplitSizeActual,
586	- layoutOTmp, layoutO, layoutUpdate,	609	+ rescaleOPingPongFlag, glFlag, taskPingPongFlag);
587	- actualBlockShapePV,
588	- nIdx, isLastNTile, qHeadSplitSizeActual, rescaleOPingPongFlag, glFlag, taskPingPongFlag);
589	#endif	610	#endif
590	}	611	}
591		612
	@@ -656,9 +677,8 @@ public:
656	uint32_t aivNum = AscendC::GetBlockNum() * AscendC::GetSubBlockNum();	677	uint32_t aivNum = AscendC::GetBlockNum() * AscendC::GetSubBlockNum();
657	uint32_t aivId = AscendC::GetBlockIdx();	678	uint32_t aivId = AscendC::GetBlockIdx();
658		679
659	- uint32_t headsProcess = (COMPUTE_ELE_NUM / embed) > HEADS_PROCESS_MAX	680	+ uint32_t headsProcess =
660	- ? HEADS_PROCESS_MAX	681	+ (COMPUTE_ELE_NUM / embed) > HEADS_PROCESS_MAX ? HEADS_PROCESS_MAX : (COMPUTE_ELE_NUM / embed);
661	- : (COMPUTE_ELE_NUM / embed);
662	uint32_t loopsPerBatch = (qHeads + headsProcess - 1) / headsProcess;	682	uint32_t loopsPerBatch = (qHeads + headsProcess - 1) / headsProcess;
663	uint32_t loopsTotal = batch * loopsPerBatch;	683	uint32_t loopsTotal = batch * loopsPerBatch;
664		684
	@@ -690,12 +710,14 @@ public:
690	actualHeads = qHeads - loopIdxInBatch * headsProcess;	710	actualHeads = qHeads - loopIdxInBatch * headsProcess;
691	}	711	}
692		712
693	- for(uint32_t qSeqIdx = 0;qSeqIdx < qSeqlen;qSeqIdx++){	713	+ for (uint32_t qSeqIdx = 0; qSeqIdx < qSeqlen; qSeqIdx++) {
694	epilogueMLAFDRescaleO(	714	epilogueMLAFDRescaleO(
695	gO[oAddr + qSeqIdx * qHeads * embed + loopIdxInBatch * headsProcess * embed],	715	gO[oAddr + qSeqIdx * qHeads * embed + loopIdxInBatch * headsProcess * embed],
696	- gOCoreTmp[oFdOffset * maxKvSplitCoreNum + qSeqIdx * qHeads * embed * maxKvSplitCoreNum +	716	+ gOCoreTmp
697	- loopIdxInBatch * headsProcess * maxKvSplitCoreNum * embed],	717	+ [oFdOffset * maxKvSplitCoreNum + qSeqIdx * qHeads * embed * maxKvSplitCoreNum +
698	- gl[lOffset + qSeqIdx * qHeads * maxKvSplitCoreNum + loopIdxInBatch * headsProcess * maxKvSplitCoreNum],	718	+ loopIdxInBatch * headsProcess * maxKvSplitCoreNum * embed],
		719	+ gl[lOffset + qSeqIdx * qHeads * maxKvSplitCoreNum +
		720	+ loopIdxInBatch * headsProcess * maxKvSplitCoreNum],
699	actualHeads, headsProcess, embed, kvSplitCoreNum);	721	actualHeads, headsProcess, embed, kvSplitCoreNum);
700	}	722	}
701	}	723	}
	@@ -711,20 +733,9 @@ private:
711	};	733	};
712		734
713	template <class Dtype>	735	template <class Dtype>
714	-CATLASS_GLOBAL void MLA(uint64_t hardwareSyncAddr,	736	+CATLASS_GLOBAL void MLA(
715	- GM_ADDR q,	737	+ uint64_t hardwareSyncAddr, GM_ADDR q, GM_ADDR qRope, GM_ADDR k, GM_ADDR kRope, GM_ADDR blockTables, GM_ADDR o,
716	- GM_ADDR qRope,	738	+ GM_ADDR s, GM_ADDR p, GM_ADDR oTmp, GM_ADDR oUpdate, GM_ADDR oCoreTmp, GM_ADDR l, GM_ADDR tiling)
717	- GM_ADDR k,
718	- GM_ADDR kRope,
719	- GM_ADDR blockTables,
720	- GM_ADDR o,
721	- GM_ADDR s,
722	- GM_ADDR p,
723	- GM_ADDR oTmp,
724	- GM_ADDR oUpdate,
725	- GM_ADDR oCoreTmp,
726	- GM_ADDR l,
727	- GM_ADDR tiling)
728	{	739	{
729	// Set hardware sync address	740	// Set hardware sync address
730	AscendC::SetSyncBaseAddr(hardwareSyncAddr);	741	AscendC::SetSyncBaseAddr(hardwareSyncAddr);
	@@ -786,8 +797,8 @@ CATLASS_GLOBAL void MLA(uint64_t hardwareSyncAddr,
786	Epilogue::Block::BlockEpilogue<Epilogue::EpilogueAtlasA2MLAFDRescaleO<ComputeEleNum>, OType, lType>;	797	Epilogue::Block::BlockEpilogue<Epilogue::EpilogueAtlasA2MLAFDRescaleO<ComputeEleNum>, OType, lType>;
787		798
788	// Kernel level	799	// Kernel level
789	- using MLAKernel = MLAKernel<BlockMmadQK, BlockMmadPV, EpilogueMLASoftmax,	800	+ using MLAKernel =
790	- EpilogueMLARescaleO, EpilogueMLAFDRescaleO>;	801	+ MLAKernel<BlockMmadQK, BlockMmadPV, EpilogueMLASoftmax, EpilogueMLARescaleO, EpilogueMLAFDRescaleO>;
791	typename MLAKernel::Params params{q, qRope, k, kRope, blockTables, o, s, p, oTmp, oUpdate, oCoreTmp, l, tiling};	802	typename MLAKernel::Params params{q, qRope, k, kRope, blockTables, o, s, p, oTmp, oUpdate, oCoreTmp, l, tiling};
792		803
793	// call kernel	804	// call kernel

Mtests/optest/kernels/19_mla/tiling/mla_tiling.cpp+76-9

	@@ -154,10 +154,17 @@ void GetTilingHead(const MLAInfo &mlaInfo, uint32_t *tilingHost, const uint32_t
154		154
155	uint32_t GetKVSplitParam(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *tilingHost, const std::vector<uint32_t>& sortedIndices)	155	uint32_t GetKVSplitParam(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *tilingHost, const std::vector<uint32_t>& sortedIndices)
156	{	156	{
157	- bool isKVSplit = (tilingHost[TILING_BATCH] <= blockDim * SPLITKV_RATION && tilingHost[TILING_MAX_QSEQLEN] == 1);	157	+ // Only split KV when the longest sequence is long enough to be worth distributing across
		158	+ // cores (>= blockDim * KV_SEQLEN_SLICE * 2). For short sequences, splitting adds extra
		159	+ // flash-decoding accumulation that degrades the max/mean relative error (MARE/MERE) on
		160	+ // near-zero output elements without any performance benefit. This matches the reference
		161	+ // implementation's split decision (blockDim = AIC core count = 20 here).
		162	+ bool isKVSplit = (tilingHost[TILING_MAX_KVSEQLEN] >= blockDim * KV_SEQLEN_SLICE * NUM2) &&
		163	+ (tilingHost[TILING_BATCH] <= blockDim * SPLITKV_RATION && tilingHost[TILING_MAX_QSEQLEN] == 1);
158		164
159	if (!isKVSplit) {	165	if (!isKVSplit) {
160	tilingHost[TILING_KVCORENUM] = 1;	166	tilingHost[TILING_KVCORENUM] = 1;
		167	+ tilingHost[TILING_PROCESSNUM] = 0;
161	for (int32_t seqIdx = 0; seqIdx < mlaInfo.batch; seqIdx++) {	168	for (int32_t seqIdx = 0; seqIdx < mlaInfo.batch; seqIdx++) {
162	uint32_t sortSeqIdx = sortedIndices[seqIdx];	169	uint32_t sortSeqIdx = sortedIndices[seqIdx];
163	uint32_t kvSeqLen = *(mlaInfo.kvSeqLen + sortSeqIdx);	170	uint32_t kvSeqLen = *(mlaInfo.kvSeqLen + sortSeqIdx);
	@@ -168,22 +175,82 @@ uint32_t GetKVSplitParam(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *t
168	return tilingHost[TILING_BATCH];	175	return tilingHost[TILING_BATCH];
169	}	176	}
170		177
171	- uint32_t MAX_KV_SPLIT_NUM = 20;	178	+ // T+1-style balanced KV split: distribute available cores among the (q==1) decode tasks
172	- for (int32_t seqIdx = 0; seqIdx < mlaInfo.batch; seqIdx++) {	179	+ // by repeatedly splitting the currently heaviest task, with a minimum tokens-per-split floor.
		180	+ // Each batch is one task here (split path requires maxQseqlen == 1).
		181	+ const uint32_t MIN_TOKENS_PER_SPLIT = static_cast<uint32_t>(mlaInfo.blockSize * NUM4);
		182	+ // Cap the per-task split count to a proven-safe value (the original common path always
		183	+ // used 8). This keeps each task's kvSplitCoreNum within the range the flash-decoding
		184	+ // reduction handles correctly, and together with totalAllocated <= blockDim guarantees
		185	+ // sum(kvSplitCoreNum) <= blockDim (i.e. at most one task per core).
		186	+ const uint32_t MAX_SPLIT_PER_TASK = 8;
		187	+ uint32_t taskCount = static_cast<uint32_t>(mlaInfo.batch);
		188	+
		189	+ std::vector<uint32_t> alignedKvLens(taskCount);
		190	+ for (uint32_t seqIdx = 0; seqIdx < taskCount; seqIdx++) {
		191	+ uint32_t sortSeqIdx = sortedIndices[seqIdx];
		192	+ uint32_t kvSeqLen = *(mlaInfo.kvSeqLen + sortSeqIdx);
		193	+ alignedKvLens[seqIdx] = RoundUp(kvSeqLen, static_cast<uint32_t>(mlaInfo.blockSize));
		194	+ }
		195	+
		196	+ std::vector<uint32_t> batchSplitNum(taskCount, 1);
		197	+ uint32_t totalAllocated = taskCount;
		198	+ while (totalAllocated < blockDim) {
		199	+ int32_t maxLoadIdx = -1;
		200	+ uint32_t maxLoad = 0;
		201	+ for (uint32_t i = 0; i < taskCount; i++) {
		202	+ uint32_t currentLoad = alignedKvLens[i] / batchSplitNum[i];
		203	+ uint32_t nextLoad = alignedKvLens[i] / (batchSplitNum[i] + 1);
		204	+ if (batchSplitNum[i] < MAX_SPLIT_PER_TASK && nextLoad >= MIN_TOKENS_PER_SPLIT && currentLoad > maxLoad) {
		205	+ maxLoad = currentLoad;
		206	+ maxLoadIdx = static_cast<int32_t>(i);
		207	+ }
		208	+ }
		209	+ if (maxLoadIdx == -1) {
		210	+ break;
		211	+ }
		212	+ batchSplitNum[maxLoadIdx]++;
		213	+ totalAllocated++;
		214	+ }
		215	+
		216	+ uint32_t MAX_KV_SPLIT_NUM = 1;
		217	+ for (uint32_t num : batchSplitNum) {
		218	+ if (num > MAX_KV_SPLIT_NUM) {
		219	+ MAX_KV_SPLIT_NUM = num;
		220	+ }
		221	+ }
		222	+
		223	+ // Build the cumulative-task prefix sum so the kernel can map process -> (task, kvIdx) tightly.
		224	+ uint32_t cuTaskVal = 0;
		225	+ int32_t cuTaskIdx = 0;
		226	+ tilingHost[CUTASK_START_OFFSET + cuTaskIdx++] = cuTaskVal;
		227	+ for (uint32_t seqIdx = 0; seqIdx < taskCount; seqIdx++) {
173	int32_t tilingOffset = seqIdx * TILING_PARA_SIZE + TILING_HEAD_SIZE;	228	int32_t tilingOffset = seqIdx * TILING_PARA_SIZE + TILING_HEAD_SIZE;
174	uint32_t sortSeqIdx = sortedIndices[seqIdx];	229	uint32_t sortSeqIdx = sortedIndices[seqIdx];
175	uint32_t kvSeqLen = *(mlaInfo.kvSeqLen + sortSeqIdx);	230	uint32_t kvSeqLen = *(mlaInfo.kvSeqLen + sortSeqIdx);
176		231
177	- uint32_t kvSeqAlign = RoundUp(kvSeqLen, static_cast<uint32_t>(mlaInfo.blockSize));	232	+ uint32_t kvSplitPerCore;
178	- uint32_t kvSeqBlockNum = kvSeqAlign / mlaInfo.blockSize;	233	+ uint32_t kvSplitCoreNum;
179	- uint32_t kvBlockPerCore = CeilDiv(kvSeqBlockNum, MAX_KV_SPLIT_NUM);	234	+ if (kvSeqLen == 0) {
180	- uint32_t kvSplitPerCore = kvBlockPerCore * mlaInfo.blockSize;	235	+ kvSplitPerCore = static_cast<uint32_t>(mlaInfo.blockSize);
181	- uint32_t kvSplitCoreNum = CeilDiv(kvSeqLen, kvSplitPerCore);	236	+ kvSplitCoreNum = 1;
		237	+ } else {
		238	+ uint32_t allocSplit = batchSplitNum[seqIdx];
		239	+ uint32_t kvSeqAlign = RoundUp(kvSeqLen, static_cast<uint32_t>(mlaInfo.blockSize));
		240	+ uint32_t kvSeqBlockNum = kvSeqAlign / mlaInfo.blockSize;
		241	+ uint32_t kvBlockPerCore = CeilDiv(kvSeqBlockNum, allocSplit);
		242	+ kvSplitPerCore = kvBlockPerCore * mlaInfo.blockSize;
		243	+ kvSplitCoreNum = CeilDiv(kvSeqLen, kvSplitPerCore);
		244	+ }
182	tilingHost[tilingOffset + NUM15] = kvSplitPerCore;	245	tilingHost[tilingOffset + NUM15] = kvSplitPerCore;
183	tilingHost[tilingOffset + NUM16] = kvSplitCoreNum;	246	tilingHost[tilingOffset + NUM16] = kvSplitCoreNum;
		247	+ cuTaskVal += kvSplitCoreNum;
		248	+ tilingHost[CUTASK_START_OFFSET + cuTaskIdx++] = cuTaskVal;
184	}	249	}
		250	+ tilingHost[TILING_PROCESSNUM] = cuTaskVal;
185	tilingHost[TILING_KVCORENUM] = MAX_KV_SPLIT_NUM;	251	tilingHost[TILING_KVCORENUM] = MAX_KV_SPLIT_NUM;
186	std::cout << "TILING_MAX_KVCORENUM = " << tilingHost[TILING_KVCORENUM] << std::endl;	252	std::cout << "TILING_MAX_KVCORENUM = " << tilingHost[TILING_KVCORENUM] << std::endl;
		253	+ std::cout << "TILING_PROCESSNUM = " << tilingHost[TILING_PROCESSNUM] << std::endl;
187		254
188	// Set lOffsetInfo and OfdOffsetInfo	255	// Set lOffsetInfo and OfdOffsetInfo
189	AddrOffsets addrOffsets;	256	AddrOffsets addrOffsets;
	@@ -201,7 +268,7 @@ uint32_t GetKVSplitParam(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *t
201	addrOffsets.addrOFdSeqOffset += static_cast<uint64_t>(mlaInfo.numHeads * qSeqlen * mlaInfo.embeddingSize);	268	addrOffsets.addrOFdSeqOffset += static_cast<uint64_t>(mlaInfo.numHeads * qSeqlen * mlaInfo.embeddingSize);
202	}	269	}
203		270
204	- return mlaInfo.batch * MAX_KV_SPLIT_NUM;	271	+ return cuTaskVal;
205	}	272	}
206		273
207	uint32_t GetKVSplitParamSpec(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *tilingHost, const std::vector<uint32_t>& sortedIndices)	274	uint32_t GetKVSplitParamSpec(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *tilingHost, const std::vector<uint32_t>& sortedIndices)