Star489
279
代码
代码
Issues74
Pull Requests49
Actions
项目讨论
Wiki
分析
Star489
279
已合并
feat: mla算子优化 — KV-split均衡分核调度和精度修复 + cpu_low.bin读取修复 #794
华工-林鑫创建于 7月1日
feat: mla算子优化 — KV-split均衡分核调度和精度修复 + cpu_low.bin读取修复 #794
已合并
华工-林鑫创建于 7月1日
从xlinn/catlassmerge2合入到cann/catlassmaster
共 34 个文件变更+5628-2230
Mexamples/19_mla/README.md+5-4
@@ -7,9 +7,10 @@
7│ ├── CMakeLists.txt # CMake编译文件7│ ├── CMakeLists.txt # CMake编译文件
8│ ├── gen_data.py8│ ├── gen_data.py
9│ ├── kernel_common.hpp #两个不同的kernel实现中的共同变量与宏9│ ├── kernel_common.hpp #两个不同的kernel实现中的共同变量与宏
10-│ ├── main.cpp10+│ ├── mla.cpp
11│ ├── mla_kernel.cpp # MLA TP 2/4/8 模板11│ ├── mla_kernel.cpp # MLA TP 2/4/8 模板
12│ ├── mla_kernel_tp1_spec.cpp # MLA TP 1 模板12│ ├── mla_kernel_tp1_spec.cpp # MLA TP 1 模板
13+│ ├── amla_kernel_tp1_spec.cpp # AMLA TP 1 模板
13│ └── README.md14│ └── README.md
14```15```
15 16 
@@ -20,7 +21,7 @@
20 21 
21```22```
22# 在./examples/19_mla目录下执行23# 在./examples/19_mla目录下执行
23-python gen_data.py 1 1 128 16 16 128 half24+python gen_data.py 4 "1,2,3,4" "128,256,512,1024" 16 64 128 half
24# 输入参数分别对应 batchSize，qSeqlen，kvSeqlen, qheadNum，numBlock, blockSize25# 输入参数分别对应 batchSize，qSeqlen，kvSeqlen, qheadNum，numBlock, blockSize
25# qSeqlen表示需要推理的token个数，支持范围为1~4，即常规decode与mtp场景26# qSeqlen表示需要推理的token个数，支持范围为1~4，即常规decode与mtp场景
26# kvSeqlen表示输入的序列长度27# kvSeqlen表示输入的序列长度
@@ -49,9 +50,9 @@ python gen_data.py 1 1 128 16 16 128 half
49# 编译指定用例（在CATLASS代码仓目录下）50# 编译指定用例（在CATLASS代码仓目录下）
50bash scripts/build.sh 19_mla51bash scripts/build.sh 19_mla
51cd output/bin52cd output/bin
52-./19_mla 1 1 128 16 16 12853+./19_mla 4 "1,2,3,4" "128,256,512,1024" 16 64 128
53# 此处的参数和生成数据的参数保持一致54# 此处的参数和生成数据的参数保持一致
54-# 完整参数为 batchSize, qSeqlen, kvSeqlen, qheadNum, numBlock, blockSize [--dtype DTYPE --datapath DATA_PATH --device DEVICE_ID]，dtype默认为half, datapath默认为../../examples/19_mla/data, device默认为0。55+# 完整参数为 batch "qSeqlenList" "kvSeqlenList" numHeads numBlocks blockSize [--dtype DTYPE --datapath DATA_PATH --device DEVICE_ID]，dtype默认为half, datapath默认为../../examples/19_mla/data, device默认为0。
55```56```
56 57 
57执行结果如下，说明精度比对成功。58执行结果如下，说明精度比对成功。
Aexamples/19_mla/amla_kernel_tp1_spec.cpp+565-0
@@ -0,0 +1,565 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
S
Ssunhao_hw7月7日

新增文件的copyright修改为2026，统一排查下

likedislike
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include "catlass/catlass.hpp"
12+#include "catlass/arch/arch.hpp"
13+#include "catlass/layout/layout.hpp"
14+ 
15+#include "catlass/gemm/block/block_mmad.hpp"
16+#include "catlass/gemm/dispatch_policy.hpp"
17+#include "catlass/gemm/gemm_type.hpp"
18+ 
19+#include "catlass/arch/cross_core_sync.hpp"
20+#include "catlass/arch/resource.hpp"
21+#include "catlass/epilogue/block/block_epilogue.hpp"
22+#include "catlass/epilogue/dispatch_policy.hpp"
23+ 
24+#include "kernel_common.hpp"
25+ 
26+using namespace Catlass;
27+/*
28+This example demonstrates how to compute amla.
29+*/
30+template <class BlockMmadQK, class BlockMmadAMLAPV, class EpilogueAMLASoftmax, class EpilogueAMLARescaleO,
31+ class EpilogueMLAFDRescaleO>
32+class AMLAKernelTp1Spec {
33+public:
34+ using ArchTag = typename BlockMmadQK::ArchTag;
35+ using L1TileShape = typename BlockMmadQK::L1TileShape;
36+ using ElementQ = typename BlockMmadQK::ElementA;
37+ using LayoutQ = typename BlockMmadQK::LayoutA;
38+ using ElementK = typename BlockMmadQK::ElementB;
39+ using LayoutK = typename BlockMmadQK::LayoutB;
40+ using ElementS = typename BlockMmadQK::ElementC;
41+ using LayoutS = typename BlockMmadQK::LayoutC;
42+ 
43+ using ElementP = typename BlockMmadAMLAPV::ElementA;
44+ using LayoutP = typename BlockMmadAMLAPV::LayoutA;
45+ using ElementV = typename BlockMmadAMLAPV::ElementB;
46+ using LayoutV = typename BlockMmadAMLAPV::LayoutB;
47+ 
48+ using ElementMask = half;
49+ 
50+ using ElementO = typename EpilogueAMLARescaleO::ElementOutput;
51+ using LayoutO = typename EpilogueAMLARescaleO::LayoutOutput;
52+ 
53+ using ElementOTmp = typename EpilogueAMLARescaleO::ElementInput;
54+ using LayoutOTmp = typename EpilogueAMLARescaleO::LayoutInput;
55+ 
56+ using ElementUpdate = typename EpilogueAMLARescaleO::ElementUpdate;
57+ using LayoutUpdate = typename EpilogueAMLARescaleO::LayoutUpdate;
58+ 
59+ static constexpr uint32_t KV_SPLIT_MAX = EpilogueMLAFDRescaleO::KV_SPLIT_MAX;
60+ static constexpr uint32_t HEADS_PROCESS_MAX = EpilogueMLAFDRescaleO::HEADS_PROCESS_MAX;
61+ static constexpr uint32_t COMPUTE_ELE_NUM = EpilogueMLAFDRescaleO::COMPUTE_ELE_NUM;
62+ 
63+ /// Parameters structure
64+ struct Params {
65+ // Data members
66+ GM_ADDR q;
67+ GM_ADDR qRope;
68+ GM_ADDR k;
69+ GM_ADDR kRope;
70+ GM_ADDR blockTables;
71+ GM_ADDR o;
72+ GM_ADDR s;
73+ GM_ADDR p;
74+ GM_ADDR oTmp;
75+ GM_ADDR oUpdate;
76+ GM_ADDR oCoreTmp;
77+ GM_ADDR l;
78+ GM_ADDR tiling;
79+ 
80+ // Methods
81+ CATLASS_DEVICE
82+ Params() {}
83+ 
84+ CATLASS_DEVICE
85+ Params(GM_ADDR q_, GM_ADDR qRope_, GM_ADDR k_, GM_ADDR kRope_, GM_ADDR blockTables_, GM_ADDR o_, GM_ADDR s_,
86+ GM_ADDR p_, GM_ADDR oTmp_, GM_ADDR oUpdate_, GM_ADDR oCoreTmp_, GM_ADDR l_, GM_ADDR tiling_)
87+ : q(q_), qRope(qRope_), k(k_), kRope(kRope_), blockTables(blockTables_), o(o_), s(s_), p(p_), oTmp(oTmp_),
88+ oUpdate(oUpdate_), oCoreTmp(oCoreTmp_), l(l_), tiling(tiling_)
89+ {
90+ }
91+ };
92+ 
93+ // Methods
94+ CATLASS_DEVICE
95+ AMLAKernelTp1Spec() {}
96+ 
97+ CATLASS_DEVICE void operator()(Params const &params)
98+ {
99+#if defined(__DAV_CUBE__)
100+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
101+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID1);
102+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID2);
103+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID3);
104+ 
105+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
106+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID1);
107+ 
108+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
109+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID1);
110+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID2);
111+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID3);
112+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID4);
113+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID5);
114+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID6);
115+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID7);
116+#elif defined(__DAV_VEC__)
117+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID0);
118+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
119+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID1);
120+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID2);
121+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID3);
122+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID4);
123+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID5);
124+ 
125+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
126+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID1);
127+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID4);
128+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID2);
129+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID3);
130+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID5);
131+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);
132+#endif
133+ 
134+ // Get the memory offset address of the input on Global Memory
135+#if defined(__DAV_CUBE__)
136+ AscendC::GlobalTensor<ElementQ> gQ;
137+ gQ.SetGlobalBuffer((__gm__ ElementQ *)params.q);
138+ AscendC::GlobalTensor<ElementQ> gQRope;
139+ gQRope.SetGlobalBuffer((__gm__ ElementQ *)params.qRope);
140+ AscendC::GlobalTensor<ElementK> gK;
141+ gK.SetGlobalBuffer((__gm__ ElementK *)params.k);
142+ AscendC::GlobalTensor<ElementK> gKRope;
143+ gKRope.SetGlobalBuffer((__gm__ ElementK *)params.kRope);
144+ AscendC::GlobalTensor<int32_t> gblockTable;
145+ gblockTable.SetGlobalBuffer((__gm__ int32_t *)(params.blockTables));
146+#elif defined(__DAV_VEC__)
147+ AscendC::GlobalTensor<ElementO> gO;
148+ gO.SetGlobalBuffer((__gm__ ElementO *)params.o);
149+ AscendC::GlobalTensor<ElementOTmp> gOUpdate;
150+ gOUpdate.SetGlobalBuffer((__gm__ ElementOTmp *)params.oUpdate);
151+ AscendC::GlobalTensor<ElementOTmp> gOCoreTmp;
152+ gOCoreTmp.SetGlobalBuffer((__gm__ ElementOTmp *)params.oCoreTmp);
153+ AscendC::GlobalTensor<ElementOTmp> gl;
154+ gl.SetGlobalBuffer((__gm__ ElementOTmp *)params.l);
155+ AscendC::GlobalTensor<float> gTilingFp64;
156+ gTilingFp64.SetGlobalBuffer((__gm__ float *)params.tiling);
157+#endif
158+ AscendC::GlobalTensor<ElementS> gS;
159+ gS.SetGlobalBuffer((__gm__ ElementS *)params.s);
160+ AscendC::GlobalTensor<ElementP> gP;
161+ gP.SetGlobalBuffer((__gm__ ElementP *)params.p);
162+ AscendC::GlobalTensor<ElementOTmp> gOTmp;
163+ gOTmp.SetGlobalBuffer((__gm__ ElementOTmp *)params.oTmp);
164+#if defined(__DAV_VEC__)
165+ AscendC::GlobalTensor<int32_t> &gOTmpInt32 = reinterpret_cast<AscendC::GlobalTensor<int32_t> &>(gOTmp);
166+#endif
167+ AscendC::GlobalTensor<uint32_t> gTiling;
168+ gTiling.SetGlobalBuffer((__gm__ uint32_t *)params.tiling);
169+ 
170+#if defined(__DAV_CUBE__)
171+ uint32_t coreIdx = AscendC::GetBlockIdx();
172+#elif defined(__DAV_VEC__)
173+ uint32_t coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum();
174+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
175+ float tor = gTilingFp64.GetValue(TILING_TOR);
176+#endif
177+ uint32_t coreNum = AscendC::GetBlockNum();
178+ 
179+ // Get tiling parameters
180+ uint32_t batch = gTiling.GetValue(TILING_BATCH);
181+ uint32_t qHeads = gTiling.GetValue(TILING_NUMHEADS);
182+ uint32_t blockSize = gTiling.GetValue(TILING_BLOCKSIZE);
183+ uint32_t maxNumBlocksPerQuery = gTiling.GetValue(TILING_MAXBLOCKS);
184+ uint32_t totalTaskNumSpec = gTiling.GetValue(TILING_TOTAL_QTOKENS);
185+ uint32_t tilingHeadSize = gTiling.GetValue(TILING_HEADSIZE);
186+ uint32_t tilingParaSize = gTiling.GetValue(TILING_PARASIZE);
187+ uint32_t maxKvSplitCoreNum = gTiling.GetValue(TILING_KVCORENUM);
188+ uint32_t formerTaskNum = gTiling.GetValue(TILING_FORMERTASKNUM);
189+ uint32_t tailTaskNum = gTiling.GetValue(TILING_TAILTASKNUM);
190+ 
191+ uint32_t embed = NUM512;
192+ uint32_t embedRope = NUM64;
193+ uint32_t kvHeads = NUM1;
194+ uint32_t strideQO = qHeads * embed;
195+ uint32_t strideQORope = qHeads * embedRope;
196+ uint32_t embedRound = RoundUp<BLOCK_SIZE>(embed);
197+ uint32_t glFlag[2] = {1, 1};
198+#if defined(__DAV_CUBE__)
199+ BlockMmadQK blockMmadQK(resource);
200+ BlockMmadAMLAPV blockMmadAMLAPV(resource);
201+#elif defined(__DAV_VEC__)
202+ EpilogueAMLASoftmax epilogueAMLATP1Softmax(resource, tor, maxKvSplitCoreNum);
203+ EpilogueAMLARescaleO epilogueAMLATP1RescaleO(resource, maxKvSplitCoreNum);
204+#endif
205+ 
206+ uint32_t pingpongIdx = 0;
207+ uint32_t pvLoopPingpongIdx = 0;
208+ uint32_t qkInLoopPingpongIdx = 0;
209+ uint32_t qkLoopPingpongIdx = 0;
210+ bool isForward = true;
211+ uint32_t taskPingPongFlag = 0;
212+ for (uint32_t process = coreIdx; process < formerTaskNum; process += uint32_t(coreNum)) {
213+ uint32_t bigProcess = process - (process % coreNum) + (coreNum - 1);
214+ bigProcess = (bigProcess > formerTaskNum - 1) ? (formerTaskNum - 1) : bigProcess;
215+ uint32_t realProcess = isForward ? process : (bigProcess - process % uint32_t(coreNum));
216+ isForward = !isForward;
217+ 
218+ // Get the offset of each core on the GM
219+ uint32_t offsetTiling = tilingHeadSize + tilingParaSize * realProcess;
220+ uint32_t curBatch = gTiling.GetValue(offsetTiling);
221+ uint32_t curTokenWiseOffset = gTiling.GetValue(offsetTiling + 1);
222+ uint32_t kvSeqlen = gTiling.GetValue(offsetTiling + 2);
223+ uint64_t gmOffsetQ = (uint64_t)(curTokenWiseOffset * strideQO);
224+ uint64_t gmOffsetQRope = (uint64_t)(curTokenWiseOffset * strideQORope);
225+ uint64_t gmOffsetO = curTokenWiseOffset * qHeads * embed;
226+ 
227+ if (kvSeqlen == 0) {
228+ continue;
229+ }
230+ uint32_t nLoop = (kvSeqlen + blockSize - 1) / blockSize;
231+ uint32_t stackSeqTile = blockSize * UNIT_BLOCK_STACK_NUM;
232+ uint32_t rowNum = qHeads;
233+ uint32_t rowNumRound = RoundUp<BLOCK_SIZE>(rowNum);
234+ uint64_t gmOffsetBlockTable = curBatch * maxNumBlocksPerQuery;
235+ 
236+ if (process == coreIdx) {
237+ uint32_t nIdx = 0;
238+ if (nIdx + UNIT_BLOCK_STACK_NUM > nLoop - 1) {
239+ stackSeqTile = kvSeqlen - nIdx * blockSize;
240+ } else {
241+ stackSeqTile = blockSize * UNIT_BLOCK_STACK_NUM;
242+ }
243+ uint32_t stackSeqTileRound = RoundUp<BLOCK_SIZE>(stackSeqTile);
244+ uint32_t gSPingPongFlag = (pingpongIdx / UNIT_BLOCK_STACK_NUM) % 2;
245+#if defined(__DAV_CUBE__)
246+ LayoutQ layoutQ(rowNum, embed);
247+ LayoutQ layoutQRope(rowNum, embedRope);
248+ LayoutK layoutK(embed, stackSeqTile);
249+ LayoutK layoutKRope(embedRope, stackSeqTile);
250+ LayoutS layoutS(rowNumRound, stackSeqTileRound);
251+ GemmCoord actualBlockShapeQK{rowNum, stackSeqTile, embed + embedRope};
252+ uint64_t gmOffseS =
253+ (uint64_t)coreIdx * TMP_SIZE_DECODER * 4 + (uint64_t)gSPingPongFlag * TMP_SIZE_DECODER * 2;
254+ // Calculate Q * K^T
255+ blockMmadQK(gQ[gmOffsetQ], gQRope[gmOffsetQRope], gK, gKRope, gblockTable[gmOffsetBlockTable],
256+ gS[gmOffseS], layoutQ, layoutQRope, layoutK, layoutKRope, layoutS, actualBlockShapeQK,
257+ nIdx, nLoop, blockSize, kvSeqlen, qkInLoopPingpongIdx, qkLoopPingpongIdx);
258+ Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady);
259+#elif defined(__DAV_VEC__)
260+ LayoutP layoutP(rowNum, stackSeqTile, stackSeqTileRound);
261+ LayoutS layoutS(rowNum, stackSeqTile, stackSeqTileRound);
262+ LayoutOTmp layoutOTmp(rowNum, embed, embedRound);
263+ GemmCoord actualBlockShapeQK{rowNum, stackSeqTile, embed};
264+ uint32_t gmOffsetP = (uint64_t)coreIdx * TMP_SIZE * 2 +
265+ (uint64_t)subBlockIdx * rowNum / 2 * stackSeqTileRound +
266+ (uint64_t)gSPingPongFlag * TMP_SIZE;
267+ uint32_t gmOffsetS = (int64_t)coreIdx * TMP_SIZE_DECODER * 4 +
268+ (int64_t)subBlockIdx * rowNum / 2 * stackSeqTileRound +
269+ (uint64_t)gSPingPongFlag * TMP_SIZE_DECODER * 2;
270+ uint64_t gmOffsetOTmp = (uint64_t)(coreIdx * TMP_SIZE * 2 +
271+ subBlockIdx * rowNum * embedRound / 2 +
272+ taskPingPongFlag * TMP_SIZE);
273+ uint32_t isLastNTile = (nIdx + UNIT_BLOCK_STACK_NUM > nLoop - 1) ? 1 : 0;
274+ // Softmax one-stage calculation
275+ epilogueAMLATP1Softmax(gP[gmOffsetP], gS[gmOffsetS], gOTmpInt32[gmOffsetOTmp], layoutP, layoutS,
276+ layoutOTmp, actualBlockShapeQK, nIdx, isLastNTile, glFlag, taskPingPongFlag,
277+ gSPingPongFlag);
278+#endif
279+ pingpongIdx += UNIT_BLOCK_STACK_NUM;
280+ }
281+ 
282+ // Split k seqlen
283+ for (uint32_t nIdx = UNIT_BLOCK_STACK_NUM; nIdx < nLoop + UNIT_BLOCK_STACK_NUM;
284+ nIdx += UNIT_BLOCK_STACK_NUM, pingpongIdx += UNIT_BLOCK_STACK_NUM) {
285+ if (nIdx < nLoop) {
286+ if (nIdx + UNIT_BLOCK_STACK_NUM > nLoop - 1) {
287+ stackSeqTile = kvSeqlen - nIdx * blockSize;
288+ } else {
289+ stackSeqTile = blockSize * UNIT_BLOCK_STACK_NUM;
290+ }
291+ uint32_t stackSeqTileRound = RoundUp<BLOCK_SIZE>(stackSeqTile);
292+ uint32_t gSPingPongFlag = (pingpongIdx / UNIT_BLOCK_STACK_NUM) % 2;
293+#if defined(__DAV_CUBE__)
294+ // Calculate Q * K^T
295+ LayoutQ layoutQ(rowNum, embed);
296+ LayoutQ layoutQRope(rowNum, embedRope);
297+ LayoutK layoutK(embed, stackSeqTile);
298+ LayoutK layoutKRope(embedRope, stackSeqTile);
299+ LayoutS layoutS(rowNumRound, stackSeqTileRound);
300+ GemmCoord actualBlockShapeQK{rowNum, stackSeqTile, embed + embedRope};
301+ uint64_t gmOffseS =
302+ (uint64_t)coreIdx * TMP_SIZE_DECODER * 4 + (uint64_t)gSPingPongFlag * TMP_SIZE_DECODER * 2;
303+ blockMmadQK(gQ[gmOffsetQ], gQRope[gmOffsetQRope], gK, gKRope, gblockTable[gmOffsetBlockTable],
304+ gS[gmOffseS], layoutQ, layoutQRope, layoutK, layoutKRope, layoutS, actualBlockShapeQK,
305+ nIdx, nLoop, blockSize, kvSeqlen, qkInLoopPingpongIdx, qkLoopPingpongIdx);
306+ Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady);
307+#elif defined(__DAV_VEC__)
308+ LayoutP layoutP(rowNum, stackSeqTile, stackSeqTileRound);
309+ LayoutS layoutS(rowNum, stackSeqTile, stackSeqTileRound);
310+ LayoutOTmp layoutOTmp(rowNum, embed, embedRound);
311+ GemmCoord actualBlockShapeQK{rowNum, stackSeqTile, embed};
312+ uint32_t gmOffsetP = (uint64_t)coreIdx * TMP_SIZE * 2 +
313+ (uint64_t)subBlockIdx * rowNum / 2 * stackSeqTileRound +
314+ gSPingPongFlag * TMP_SIZE;
315+ uint32_t gmOffsetS = (int64_t)coreIdx * TMP_SIZE_DECODER * 4 +
316+ (int64_t)subBlockIdx * rowNum / 2 * stackSeqTileRound +
317+ gSPingPongFlag * TMP_SIZE_DECODER * 2;
318+ uint64_t gmOffsetOTmp = (uint64_t)(coreIdx * TMP_SIZE * 2 +
319+ subBlockIdx * rowNum * embedRound / 2 +
320+ taskPingPongFlag * TMP_SIZE);
321+ uint32_t isLastNTile = (nIdx + UNIT_BLOCK_STACK_NUM > nLoop - 1) ? 1 : 0;
322+ // Softmax one-stage calculation
323+ epilogueAMLATP1Softmax(gP[gmOffsetP], gS[gmOffsetS], gOTmpInt32[gmOffsetOTmp], layoutP, layoutS,
324+ layoutOTmp, actualBlockShapeQK, nIdx, isLastNTile, glFlag, taskPingPongFlag,
325+ gSPingPongFlag);
326+#endif
327+ }
328+ uint32_t nextProcess = process + uint32_t(coreNum);
329+ if (nIdx + UNIT_BLOCK_STACK_NUM > nLoop + UNIT_BLOCK_STACK_NUM - 1 && nextProcess < formerTaskNum) {
330+ bool nextIsForward = isForward;
331+ uint32_t nextBigProcess = nextProcess - (nextProcess % coreNum) + (coreNum - 1);
332+ nextBigProcess = (nextBigProcess > formerTaskNum - 1) ? (formerTaskNum - 1) : nextBigProcess;
333+ uint32_t nextRealProcess =
334+ nextIsForward ? nextProcess : (nextBigProcess - nextProcess % uint32_t(coreNum));
335+ 
336+ // Get the offset of each core on the GM
337+ uint32_t nextOffsetTiling = tilingHeadSize + tilingParaSize * nextRealProcess;
338+ uint32_t nextCurBatch = gTiling.GetValue(nextOffsetTiling);
339+ uint32_t nextCurTokenWiseOffset = gTiling.GetValue(nextOffsetTiling + 1);
340+ uint32_t nextKvSeqlen = gTiling.GetValue(nextOffsetTiling + 2);
341+ uint64_t nextGMOffsetQ = (uint64_t)(nextCurTokenWiseOffset * strideQO);
342+ uint64_t nextGMOffsetQRope = (uint64_t)(nextCurTokenWiseOffset * strideQORope);
343+ uint32_t nextNLoop = (nextKvSeqlen + blockSize - 1) / blockSize;
344+ uint64_t nextGMOffsetBlockTable = nextCurBatch * maxNumBlocksPerQuery;
345+ uint32_t nextNIdx = 0;
346+ if (nextNIdx + UNIT_BLOCK_STACK_NUM > nextNLoop - 1) {
347+ stackSeqTile = nextKvSeqlen - nextNIdx * blockSize;
348+ } else {
349+ stackSeqTile = blockSize * UNIT_BLOCK_STACK_NUM;
350+ }
351+ uint32_t stackSeqTileRound = RoundUp<BLOCK_SIZE>(stackSeqTile);
352+ uint32_t gSPingPongFlag = (pingpongIdx / UNIT_BLOCK_STACK_NUM) % 2;
353+#if defined(__DAV_CUBE__)
354+ LayoutQ layoutQ(rowNum, embed);
355+ LayoutQ layoutQRope(rowNum, embedRope);
356+ LayoutK layoutK(embed, stackSeqTile);
357+ LayoutK layoutKRope(embedRope, stackSeqTile);
358+ LayoutS layoutS(rowNumRound, stackSeqTileRound);
359+ GemmCoord nextActualBlockShapeQK{rowNum, stackSeqTile, embed + embedRope};
360+ uint64_t nextGMOffseS =
361+ (uint64_t)coreIdx * TMP_SIZE_DECODER * 4 + (uint64_t)gSPingPongFlag * TMP_SIZE_DECODER * 2;
362+ // Calculate Q * K^T
363+ blockMmadQK(gQ[nextGMOffsetQ], gQRope[nextGMOffsetQRope], gK, gKRope,
364+ gblockTable[nextGMOffsetBlockTable], gS[nextGMOffseS], layoutQ, layoutQRope, layoutK,
365+ layoutKRope, layoutS, nextActualBlockShapeQK, nextNIdx, nextNLoop, blockSize,
366+ nextKvSeqlen, qkInLoopPingpongIdx, qkLoopPingpongIdx);
367+ Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady);
368+#elif defined(__DAV_VEC__)
369+ LayoutP nextLayoutP(rowNum, stackSeqTile, stackSeqTileRound);
370+ LayoutS nextLayoutS(rowNum, stackSeqTile, stackSeqTileRound);
371+ LayoutOTmp nextLayoutOTmp(rowNum, embed, embedRound);
372+ GemmCoord nextActualBlockShapeQK{rowNum, stackSeqTile, embed};
373+ uint32_t nextGMOffsetP = (uint64_t)coreIdx * TMP_SIZE * 2 +
374+ (uint64_t)subBlockIdx * rowNum / 2 * stackSeqTileRound +
375+ gSPingPongFlag * TMP_SIZE;
376+ uint32_t nextGMOffsetS = (int64_t)coreIdx * TMP_SIZE_DECODER * 4 +
377+ (int64_t)subBlockIdx * rowNum / 2 * stackSeqTileRound +
378+ gSPingPongFlag * TMP_SIZE_DECODER * 2;
379+ uint64_t gmOffsetOTmp = (uint64_t)(coreIdx * TMP_SIZE * 2 +
380+ subBlockIdx * rowNum * embedRound / 2 +
381+ (1 - taskPingPongFlag) * TMP_SIZE);
382+ uint32_t isLastNTile = (nextNIdx + UNIT_BLOCK_STACK_NUM > nextNLoop - 1) ? 1 : 0;
383+ // Softmax one-stage calculation
384+ epilogueAMLATP1Softmax(gP[nextGMOffsetP], gS[nextGMOffsetS], gOTmpInt32[gmOffsetOTmp],
385+ nextLayoutP, nextLayoutS, nextLayoutOTmp, nextActualBlockShapeQK, nextNIdx,
386+ isLastNTile, glFlag, 1 - taskPingPongFlag, gSPingPongFlag);
387+#endif
388+ }
389+ 
390+ // Wait for the four Q * K^T calculations to complete before calculating P * V
391+ if (nIdx >= UNIT_BLOCK_STACK_NUM) {
392+ if (nIdx + UNIT_BLOCK_STACK_NUM > nLoop + UNIT_BLOCK_STACK_NUM - 1) {
393+ stackSeqTile = kvSeqlen - (nIdx - UNIT_BLOCK_STACK_NUM) * blockSize;
394+ } else {
395+ stackSeqTile = blockSize * UNIT_BLOCK_STACK_NUM;
396+ }
397+ uint32_t stackSeqTileRound = RoundUp<BLOCK_SIZE>(stackSeqTile);
398+#if defined(__DAV_CUBE__)
399+ LayoutP layoutP(rowNum, stackSeqTile, stackSeqTileRound);
400+ LayoutV layoutV(stackSeqTile, embed);
401+ LayoutOTmp layoutOTmp(rowNumRound, embedRound);
402+ GemmCoord actualBlockShapePV{rowNum, embed, stackSeqTile};
403+ uint32_t gPPingPongFlag = (pingpongIdx / UNIT_BLOCK_STACK_NUM - 1) % 2;
404+ uint64_t gmOffseP = (uint64_t)coreIdx * TMP_SIZE * 2 + (uint64_t)gPPingPongFlag * TMP_SIZE;
405+ uint64_t gmOffseOtmp = (uint64_t)(coreIdx * TMP_SIZE * 2 + taskPingPongFlag * TMP_SIZE);
406+ // Calculate P * V
407+ blockMmadAMLAPV(gP[gmOffseP], gK, gblockTable[gmOffsetBlockTable], gOTmp[gmOffseOtmp], layoutP,
408+ layoutV, layoutOTmp, actualBlockShapePV, nIdx, gPPingPongFlag, nLoop, blockSize,
409+ kvSeqlen, softmaxReady, pvLoopPingpongIdx, softmaxAutoAddReady);
410+ if (nIdx + UNIT_BLOCK_STACK_NUM > nLoop + UNIT_BLOCK_STACK_NUM - 1) {
411+ Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(pvReady);
412+ } else {
413+ Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(preAutoAddReady);
414+ }
415+#elif defined(__DAV_VEC__)
416+ if (nIdx + UNIT_BLOCK_STACK_NUM > nLoop + UNIT_BLOCK_STACK_NUM - 1) {
417+ // Wait for P * V calculation to complete
418+ uint32_t rescaleOPingPongFlag = (pingpongIdx / UNIT_BLOCK_STACK_NUM - 1) % 2;
419+ Arch::CrossCoreWaitFlag(pvReady);
420+ LayoutO layoutO(rowNum, embed);
421+ LayoutOTmp layoutOTmp(rowNum, embed, embedRound);
422+ LayoutUpdate layoutUpdate(rowNum, embed, embedRound);
423+ GemmCoord actualBlockShapePV{rowNum, embed, stackSeqTile};
424+ uint64_t gmOffsetOTmp = (uint64_t)(coreIdx * TMP_SIZE * 2 + taskPingPongFlag * TMP_SIZE);
425+ uint64_t gmOffsetUpdate = (uint64_t)(coreIdx * TMP_SIZE);
426+ // Softmax two-stage update
427+ epilogueAMLATP1RescaleO(gOTmp[gmOffsetOTmp], gOUpdate[gmOffsetUpdate], gO[gmOffsetO],
428+ gOCoreTmp[0], gl[0], layoutOTmp, layoutUpdate, layoutO,
429+ actualBlockShapePV, nIdx, rescaleOPingPongFlag, glFlag,
430+ taskPingPongFlag);
431+ }
432+#endif
433+ }
434+ }
435+ taskPingPongFlag = 1 - taskPingPongFlag;
436+ }
437+ 
438+#if defined(__DAV_CUBE__)
439+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
440+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID1);
441+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID2);
442+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID3);
443+ 
444+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
445+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID1);
446+ 
447+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
448+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID1);
449+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID2);
450+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID3);
451+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID4);
452+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID5);
453+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID6);
454+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID7);
455+#elif defined(__DAV_VEC__)
456+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID0);
457+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
458+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID1);
459+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID2);
460+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID3);
461+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID4);
462+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID5);
463+ 
464+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID4);
465+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
466+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID1);
467+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID2);
468+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID3);
469+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID5);
470+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);
471+#endif
472+ }
473+ 
474+private:
475+ Arch::Resource<ArchTag> resource;
476+ Arch::CrossCoreFlag qkReady{QK_READY_ID};
477+ Arch::CrossCoreFlag softmaxReady{SOFTMAX_READY_ID};
478+ Arch::CrossCoreFlag pvReady{PV_READY_ID};
479+ Arch::CrossCoreFlag preAutoAddReady{PRE_AUTOADD_READY_ID};
480+ Arch::CrossCoreFlag softmaxAutoAddReady{SM_AUTOADD_READY_ID};
481+};
482+ 
483+template <class Dtype>
484+CATLASS_GLOBAL void AMLATp1Spec(uint64_t hardwareSyncAddr,
485+ GM_ADDR q,
486+ GM_ADDR qRope,
487+ GM_ADDR k,
488+ GM_ADDR kRope,
489+ GM_ADDR blockTables,
490+ GM_ADDR o,
491+ GM_ADDR s,
492+ GM_ADDR p,
493+ GM_ADDR oTmp,
494+ GM_ADDR oUpdate,
495+ GM_ADDR oCoreTmp,
496+ GM_ADDR l,
497+ GM_ADDR tiling)
498+{
499+ // Set hardware sync address
500+ AscendC::SetSyncBaseAddr(hardwareSyncAddr);
501+ 
502+ using ArchTag = Arch::AtlasA2;
503+ using ElementQ = Dtype;
504+ using LayoutQ = layout::RowMajor;
505+ using ElementK = Dtype;
506+ using LayoutK = layout::ColumnMajor;
507+ using ElementV = Dtype;
508+ using LayoutV = layout::RowMajor;
509+ using ElementS = float;
510+ using LayoutS = layout::RowMajor;
511+ using ElementP = Dtype;
512+ using LayoutP = layout::RowMajor;
513+ using ElementO = Dtype;
514+ using LayoutO = layout::RowMajor;
515+ using ElementMask = Dtype;
516+ using LayoutMask = layout::RowMajor;
517+ using ElementOTmp = float;
518+ using LayoutOTmp = layout::RowMajor;
519+ using ElementUpdate = float;
520+ using LayoutUpdate = layout::RowMajor;
521+ 
522+ // L1TileShape::K must be embdding
523+ using L1TileShape = GemmShape<128, 128, 576>;
524+ using L0TileShape = L1TileShape;
525+ 
526+ // GEMM Block模块，实现Flash MLA的Q * K^T
527+ using DispatchPolicyQK = Gemm::MmadAtlasA2MLAQKTp1Spec;
528+ using QType = Gemm::GemmType<ElementQ, LayoutQ>;
529+ using KType = Gemm::GemmType<ElementK, LayoutK>;
530+ using SType = Gemm::GemmType<ElementS, LayoutS>;
531+ using BlockMmadQK = Gemm::Block::BlockMmad<DispatchPolicyQK, L1TileShape, L0TileShape, QType, KType, SType>;
532+ 
533+ // Epilogue Block模块，实现Flash MLA中当前S基块的softmax
534+ using PType = Gemm::GemmType<ElementP, LayoutP>;
535+ using MaskType = Gemm::GemmType<ElementMask, LayoutMask>;
536+ using OTmpType = Gemm::GemmType<ElementOTmp, LayoutOTmp>;
537+ using EpilogueAMLASoftmax =
538+ Epilogue::Block::BlockEpilogue<Epilogue::EpilogueAtlasA2AMLATP1Softmax, PType, SType, OTmpType, MaskType>;
539+ 
540+ // GEMM Block模块，实现Flash MLA的P * V
541+ using VType = Gemm::GemmType<ElementV, LayoutV>;
542+ using DispatchPolicyAMLAPV = Gemm::MmadAtlasA2AMLAPVTp1Spec;
543+ using BlockMmadAMLAPV = Gemm::Block::BlockMmad<DispatchPolicyAMLAPV, L1TileShape, L0TileShape, PType, VType, OTmpType>;
544+ 
545+ // Epilogue Block模块，实现Flash MLA中当前O基块的更新
546+ using OType = Gemm::GemmType<ElementO, LayoutO>;
547+ using OUpdateType = Gemm::GemmType<ElementUpdate, LayoutUpdate>;
548+ using EpilogueAMLARescaleO =
549+ Epilogue::Block::BlockEpilogue<Epilogue::EpilogueAtlasA2AMLATP1RescaleO, OType, OUpdateType, OTmpType>;
550+ 
551+ // Epilogue Block模块，实现Flash MLA中flash decoding
552+ using lType = Gemm::GemmType<ElementUpdate, LayoutUpdate>;
553+ constexpr uint32_t ComputeEleNum = 6144;
554+ using EpilogueMLAFDRescaleO =
555+ Epilogue::Block::BlockEpilogue<Epilogue::EpilogueAtlasA2MLAFDRescaleO<ComputeEleNum>, OType, lType>;
556+ 
557+ // Kernel level
558+ using MLAKernel = AMLAKernelTp1Spec<BlockMmadQK, BlockMmadAMLAPV, EpilogueAMLASoftmax,
559+ EpilogueAMLARescaleO, EpilogueMLAFDRescaleO>;
560+ typename MLAKernel::Params params{q, qRope, k, kRope, blockTables, o, s, p, oTmp, oUpdate, oCoreTmp, l, tiling};
561+ 
562+ // call kernel
563+ MLAKernel mla;
564+ mla(params);
565+}
Mexamples/19_mla/gen_data.py+45-15
@@ -49,8 +49,14 @@ class TestPagedMLAttention():
49 dtype: any49 dtype: any
50 50 
51 @classmethod51 @classmethod
52- def check_attr(cls, batch: int, q_seqlen: int, kv_seqlen: int, num_blocks: int, block_size: int):52+ def check_attr(cls, batch: int, q_seqlen_list: list, k_seqlen_list: list, num_blocks: int, block_size: int):
53- if batch * kv_seqlen > num_blocks * block_size:53+ # 检查列表长度是否与 batch size 匹配
54+ if len(q_seqlen_list) != batch or len(k_seqlen_list) != batch:
55+ logging("[ERROR] The length of q_seqlen_list and k_seqlen_list must be equal to batch size.")
56+ sys.exit()
57+ 
58+ # 检查缓存大小是否足够
59+ if sum(k_seqlen_list) > num_blocks * block_size:
54 logging("[ERROR] the number of K and V tokens is too big to fit in the paged cache.")60 logging("[ERROR] the number of K and V tokens is too big to fit in the paged cache.")
55 sys.exit()61 sys.exit()
56 62 
@@ -58,9 +64,18 @@ class TestPagedMLAttention():
58 logging("[ERROR] blockSize != 128 is not supported.")64 logging("[ERROR] blockSize != 128 is not supported.")
59 sys.exit()65 sys.exit()
60 66 
61- if q_seqlen > 4:67+ # 检查每个 q_seqlen 是否在有效范围内 [1, 4]
62- logging("[ERROR] q_seqlen > 4 is not supported.")68+ for q_seqlen in q_seqlen_list:
63- sys.exit()69+ if q_seqlen > 4 or q_seqlen < 1:
70+ logging(f"[ERROR] q_seqlen value {q_seqlen} is not in the valid range of [1, 4].")
71+ sys.exit()
72+ 
73+ # 检查每个 kv_seqlen 是否在有效范围内 [128, 16384]
74+ for kv_seqlen in k_seqlen_list:
75+ if kv_seqlen > 16384 or kv_seqlen < 128:
76+ logging(f"[ERROR] kv_seqlen value {kv_seqlen} is not in the valid range of [128, 16384].")
77+ sys.exit()
78+ 
64 79 
65 @classmethod80 @classmethod
66 def group_matmul(cls, head, kv_head, left, right):81 def group_matmul(cls, head, kv_head, left, right):
@@ -101,7 +116,7 @@ class TestPagedMLAttention():
101 sim_high = sim_high + (116 sim_high = sim_high + (
102 mask[:sim_high.shape[-2], :sim_high.shape[-1]] * self.post_mask_factor117 mask[:sim_high.shape[-2], :sim_high.shape[-1]] * self.post_mask_factor
103 ).astype(np.float32)118 ).astype(np.float32)
104- 119+ 
105 # softmax120 # softmax
106 p_high = self.softmax_numpy(sim_high)121 p_high = self.softmax_numpy(sim_high)
107 p = p_high.astype(query.dtype)122 p = p_high.astype(query.dtype)
@@ -188,7 +203,7 @@ class TestPagedMLAttention():
188 203 
189 pre_mask_factor = -10000.0204 pre_mask_factor = -10000.0
190 if gen_data_params.mask_type == 1:205 if gen_data_params.mask_type == 1:
191- mask = np.zeros(shape=(num_tokens, max_k_seqlen)).astype(np.float16)206+ mask = np.zeros(shape=(num_tokens, max_k_seqlen)).astype(gen_data_params.dtype)
192 pre_qseqlen = 0207 pre_qseqlen = 0
193 for i in range(batch_size):208 for i in range(batch_size):
194 qseqlen = gen_data_params.q_seqlen_list[i]209 qseqlen = gen_data_params.q_seqlen_list[i]
@@ -226,20 +241,35 @@ class TestPagedMLAttention():
226 os.path.join(WORKSPACE, "data", "kv_seqlen.bin"))241 os.path.join(WORKSPACE, "data", "kv_seqlen.bin"))
227 if mask:242 if mask:
228 mask.tofile(os.path.join(WORKSPACE, "data", "mask.bin"))243 mask.tofile(os.path.join(WORKSPACE, "data", "mask.bin"))
229- ref_output.astype(np.float32).tofile(os.path.join(WORKSPACE, "data", "golden.bin"))244+ ref_output.astype(gen_data_params.dtype).tofile(os.path.join(WORKSPACE, "data", "cpu_low.bin"))
245+ true_out.astype(np.float32).tofile(os.path.join(WORKSPACE, "data", "golden.bin"))
230 246 
231 247 
232if __name__ == "__main__":248if __name__ == "__main__":
233 os.makedirs(os.path.join(WORKSPACE, "data"), exist_ok=True)249 os.makedirs(os.path.join(WORKSPACE, "data"), exist_ok=True)
234 250 
251+ 
252+ # 修改命令行参数解析逻辑，以接受列表形式的输入
253+ if len(sys.argv) != 8:
254+ print("Usage: python gen_data.py <batchSize> \"<qSeqlen_list>\" \"<kvSeqlen_list>\" <qheadNum> <numBlock> <blockSize> <dtype>")
255+ print("Example: python gen_data.py 4 \"1,2,3,4\" \"128,256,512,1024\" 16 16 128 half")
256+ sys.exit(1)
257+ 
235 batch = int(sys.argv[1])258 batch = int(sys.argv[1])
236- q_seqlen = int(sys.argv[2])259+ 
237- kv_seqlen = int(sys.argv[3])260+ # 将逗号分隔的字符串解析为整数列表
261+ q_seqlen_list_str = sys.argv[2]
262+ kv_seqlen_list_str = sys.argv[3]
263+ q_seqlen_list = [int(x.strip()) for x in q_seqlen_list_str.split(',')]
264+ kv_seqlen_list = [int(x.strip()) for x in kv_seqlen_list_str.split(',')]
265+ 
238 num_head = int(sys.argv[4])266 num_head = int(sys.argv[4])
239 num_blocks = int(sys.argv[5])267 num_blocks = int(sys.argv[5])
240 block_size = int(sys.argv[6])268 block_size = int(sys.argv[6])
241 str_dtype = str(sys.argv[7])269 str_dtype = str(sys.argv[7])
242- max_kv_seqlen = kv_seqlen270+ max_kv_seqlen = max(kv_seqlen_list) if kv_seqlen_list else 0
271+ 
272+ 
243 mask_type = 0273 mask_type = 0
244 kv_heads = 1274 kv_heads = 1
245 embedding_size = 512275 embedding_size = 512
@@ -251,13 +281,13 @@ if __name__ == "__main__":
251 else:281 else:
252 logging("[ERROR] dtype must be half or bf16")282 logging("[ERROR] dtype must be half or bf16")
253 sys.exit()283 sys.exit()
254- q_seqlen_list = [q_seqlen] * batch284+ 
255- kv_seqlen_list = [kv_seqlen] * batch
256-
257 testObj = TestPagedMLAttention()285 testObj = TestPagedMLAttention()
258- testObj.check_attr(batch, q_seqlen, kv_seqlen, num_blocks, block_size)286+ 
287+ testObj.check_attr(batch, q_seqlen_list, kv_seqlen_list, num_blocks, block_size)
259 gen_data_params = testObj.GenDataParams(q_seqlen_list, kv_seqlen_list, num_head,288 gen_data_params = testObj.GenDataParams(q_seqlen_list, kv_seqlen_list, num_head,
260 kv_heads, embedding_size, embedding_size_rope,289 kv_heads, embedding_size, embedding_size_rope,
261 num_blocks, block_size, mask_type, dtype)290 num_blocks, block_size, mask_type, dtype)
291+ 
262 testObj.calc_data(gen_data_params)292 testObj.calc_data(gen_data_params)
263 293 
Mexamples/19_mla/kernel_common.hpp+5-2
@@ -8,14 +8,15 @@
8 * See LICENSE in the root of the software repository for the full text of the License.8 * See LICENSE in the root of the software repository for the full text of the License.
9 */9 */
10 10 
11+ 
11#ifndef KERNEL_COMMON12#ifndef KERNEL_COMMON
12#define KERNEL_COMMON13#define KERNEL_COMMON
13 14 
14-#include <cstdint>
15- 
16constexpr uint32_t QK_READY_ID = 1;15constexpr uint32_t QK_READY_ID = 1;
17constexpr uint32_t SOFTMAX_READY_ID = 2;16constexpr uint32_t SOFTMAX_READY_ID = 2;
18constexpr uint32_t PV_READY_ID = 3;17constexpr uint32_t PV_READY_ID = 3;
18+constexpr uint32_t PRE_AUTOADD_READY_ID = 0;
19+constexpr uint32_t SM_AUTOADD_READY_ID = 7;
19constexpr uint32_t BLOCK_SIZE = 16;20constexpr uint32_t BLOCK_SIZE = 16;
20constexpr uint32_t TMP_SIZE = 65536;21constexpr uint32_t TMP_SIZE = 65536;
21constexpr uint32_t TMP_SIZE_DECODER = 32768;22constexpr uint32_t TMP_SIZE_DECODER = 32768;
@@ -39,6 +40,8 @@ constexpr int32_t TILING_KVCORENUM = 16;
39constexpr int32_t TILING_TOTAL_QTOKENS = 18;40constexpr int32_t TILING_TOTAL_QTOKENS = 18;
40constexpr int32_t TILING_FORMERTASKNUM = 19;41constexpr int32_t TILING_FORMERTASKNUM = 19;
41constexpr int32_t TILING_TAILTASKNUM = 20;42constexpr int32_t TILING_TAILTASKNUM = 20;
43+constexpr int32_t TILING_PROCESSNUM = 21;
44+constexpr int32_t CUTASK_START_OFFSET = 25;
42constexpr int32_t TILING_BLOCKSIZE_CALC = 25;45constexpr int32_t TILING_BLOCKSIZE_CALC = 25;
43constexpr int32_t TILING_HEADDIM_K_SPLIT = 38;46constexpr int32_t TILING_HEADDIM_K_SPLIT = 38;
44constexpr int32_t TILING_HEADDIM_V_SPLIT = 39;47constexpr int32_t TILING_HEADDIM_V_SPLIT = 39;
Mexamples/19_mla/mla.cpp+101-40
@@ -16,25 +16,41 @@
16#endif16#endif
17#include <fstream>17#include <fstream>
18#include <iostream>18#include <iostream>
19+#include <sstream>
20+#include <vector>
19// Helper methods to check for errors21// Helper methods to check for errors
20#include "golden.hpp"22#include "golden.hpp"
21#include "helper.hpp"23#include "helper.hpp"
22#include "mla_kernel.cpp"24#include "mla_kernel.cpp"
23#include "mla_kernel_tp1_spec.cpp"25#include "mla_kernel_tp1_spec.cpp"
26+#include "amla_kernel_tp1_spec.cpp"
24#include "mla_tiling.h"27#include "mla_tiling.h"
25 28 
26using namespace std;29using namespace std;
27 30 
31+// Helper function to split comma-separated string into vector of uint32_t
32+vector<uint32_t> splitSeqList(const string &seqStr) {
33+ vector<uint32_t> seqList;
34+ stringstream ss(seqStr);
35+ string item;
36+ while (getline(ss, item, ',')) {
37+ seqList.push_back(static_cast<uint32_t>(stoul(item)));
38+ }
39+ return seqList;
40+}
41+ 
28// This code section describes the parameters to execute the run function.42// This code section describes the parameters to execute the run function.
29struct Options {43struct Options {
30- static constexpr auto HELPER = "Usage: mla batch qSeqlen kvSeqlen numHeads numBlocks blockSize [--dtype DTYPE "44+ static constexpr auto HELPER =
31- "--datapath DATA_PATH --device DEVICE_ID]\n";45+ "Usage: mla batch \"qSeqlenList\" \"kvSeqlenList\" numHeads numBlocks blockSize [--dtype DTYPE "
46+ "--datapath DATA_PATH --device DEVICE_ID]\n"
47+ "Example: ./19_mla 4 \"1,2,3,4\" \"128,256,512,1024\" 16 16 128\n";
32 static constexpr auto MIN_ARGS = 7;48 static constexpr auto MIN_ARGS = 7;
33 49 
34 // Define default value.50 // Define default value.
35 uint32_t batch{0};51 uint32_t batch{0};
36- uint32_t qSeqlen{0};52+ vector<uint32_t> qSeqlenList;
37- uint32_t kvSeqlen{0};53+ vector<uint32_t> kvSeqlenList;
38 uint32_t numHeads{0};54 uint32_t numHeads{0};
39 uint32_t numBlocks{0};55 uint32_t numBlocks{0};
40 uint32_t blockSize{0};56 uint32_t blockSize{0};
@@ -61,8 +77,15 @@ struct Options {
61 // Allocate arguments to parameters.77 // Allocate arguments to parameters.
62 uint32_t argIndex = 1;78 uint32_t argIndex = 1;
63 batch = atoi(argv[argIndex++]);79 batch = atoi(argv[argIndex++]);
64- qSeqlen = atoi(argv[argIndex++]);80+ string qSeqStr = argv[argIndex++];
65- kvSeqlen = atoi(argv[argIndex++]);81+ string kvSeqStr = argv[argIndex++];
82+ qSeqlenList = splitSeqList(qSeqStr);
83+ kvSeqlenList = splitSeqList(kvSeqStr);
84+ if (qSeqlenList.size() != batch || kvSeqlenList.size() != batch) {
85+ printf("Error: The number of elements in qSeqlenList and kvSeqlenList must match the batch size.\n");
86+ printf(HELPER);
87+ return -1;
88+ }
66 numHeads = atoi(argv[argIndex++]);89 numHeads = atoi(argv[argIndex++]);
67 numBlocks = atoi(argv[argIndex++]);90 numBlocks = atoi(argv[argIndex++]);
68 blockSize = atoi(argv[argIndex++]);91 blockSize = atoi(argv[argIndex++]);
@@ -109,8 +132,6 @@ static void Run(const Options &options)
109 132 
110 // Parameters initialization.133 // Parameters initialization.
111 int32_t batch = options.batch;134 int32_t batch = options.batch;
112- int32_t qSeqlen = options.qSeqlen;
113- int32_t kvSeqlen = options.kvSeqlen;
114 int32_t numHeads = options.numHeads;135 int32_t numHeads = options.numHeads;
115 int32_t kvHeads = options.kvHeads;136 int32_t kvHeads = options.kvHeads;
116 int32_t embeddingSize = options.embeddingSize;137 int32_t embeddingSize = options.embeddingSize;
@@ -120,20 +141,27 @@ static void Run(const Options &options)
120 int32_t maskType = options.maskType;141 int32_t maskType = options.maskType;
121 string dataType = options.dataType;142 string dataType = options.dataType;
122 string dataPath = options.dataPath;143 string dataPath = options.dataPath;
123- int32_t maxKvSeqlen = kvSeqlen;144+ int32_t maxKvSeqlen = 0;
145+ if (!options.kvSeqlenList.empty()) {
146+ maxKvSeqlen = options.kvSeqlenList[0];
147+ for (size_t i = 1; i < options.kvSeqlenList.size(); ++i) {
148+ if (options.kvSeqlenList[i] > maxKvSeqlen) {
149+ maxKvSeqlen = options.kvSeqlenList[i];
150+ }
151+ }
152+ }
124 153 
125 if ((dataType != "half") && (dataType != "bf16")) {154 if ((dataType != "half") && (dataType != "bf16")) {
126 cerr << "[ERROR] dtype must be 'half' or 'bf16'." << endl;155 cerr << "[ERROR] dtype must be 'half' or 'bf16'." << endl;
127 return;156 return;
128 }157 }
129 158 
130- uint32_t dTypeKey = (dataType == "half") ? 0 : 1;159+ int32_t dTypeKey = (dataType == "half") ? 0 : 1;
131- uint32_t specStraKey = (numHeads == MLATiling::NUM128) ? 1 : 0;160+ int32_t specStraKey = (numHeads == MLATiling::NUM128) ? 1 : 0;
132 161 
133 // 3 bits for tilingKey(specStraKey : 1, dTypeKey : 2)162 // 3 bits for tilingKey(specStraKey : 1, dTypeKey : 2)
134 uint32_t dTypeKeyBitLen = 2;163 uint32_t dTypeKeyBitLen = 2;
135 uint32_t tilingKey = (specStraKey << dTypeKeyBitLen) + dTypeKey;164 uint32_t tilingKey = (specStraKey << dTypeKeyBitLen) + dTypeKey;
136- std::cout << "tilingKey : " << tilingKey << std::endl;
137 165 
138 // read qNtokens num166 // read qNtokens num
139 void *qNtokens = nullptr;167 void *qNtokens = nullptr;
@@ -141,6 +169,11 @@ static void Run(const Options &options)
141 ReadFile(dataPath + "/q_ntokens.bin", qNtokens, 1 * sizeof(int32_t));169 ReadFile(dataPath + "/q_ntokens.bin", qNtokens, 1 * sizeof(int32_t));
142 int32_t numTokens = static_cast<int32_t *>(qNtokens)[0];170 int32_t numTokens = static_cast<int32_t *>(qNtokens)[0];
143 171 
172+ if ((numHeads == MLATiling::NUM128) && (numTokens % aicCoreNum <= 10) && (batch <= 40)) {
173+ tilingKey = (dTypeKey == 0) ? 7 : 8;
174+ }
175+ std::cout << "tilingKey : " << tilingKey << std::endl;
176+ 
144 // read qSeq177 // read qSeq
145 void *qSeq = nullptr;178 void *qSeq = nullptr;
146 ACL_CHECK(aclrtMallocHost(&qSeq, batch * sizeof(int32_t)));179 ACL_CHECK(aclrtMallocHost(&qSeq, batch * sizeof(int32_t)));
@@ -253,9 +286,9 @@ static void Run(const Options &options)
253 286 
254 ACL_CHECK(aclrtMemcpy(tilingDevice, tilingSize, tilingHost, tilingSize, ACL_MEMCPY_HOST_TO_DEVICE));287 ACL_CHECK(aclrtMemcpy(tilingDevice, tilingSize, tilingHost, tilingSize, ACL_MEMCPY_HOST_TO_DEVICE));
255 288 
256- uint32_t kvSplitCoreNum = *((uint32_t *)tilingHost + MLATiling::TILING_KVCORENUM);289+ uint32_t maxKvSplitCoreNum = *((uint32_t *)tilingHost + MLATiling::TILING_KVCORENUM);
257- uint64_t oFdSize = embeddingSize * numHeads * numTokens * kvSplitCoreNum * sizeof(float);290+ uint64_t oFdSize = embeddingSize * numHeads * numTokens * maxKvSplitCoreNum * sizeof(float);
258- uint64_t lSize = numTokens * numHeads * kvSplitCoreNum * sizeof(float);291+ uint64_t lSize = numTokens * numHeads * maxKvSplitCoreNum * sizeof(float);
259 292 
260 uint8_t *oCoreTmpDevice;293 uint8_t *oCoreTmpDevice;
261 ACL_CHECK(aclrtMalloc((void **)(&oCoreTmpDevice), oFdSize, ACL_MEM_MALLOC_HUGE_FIRST));294 ACL_CHECK(aclrtMalloc((void **)(&oCoreTmpDevice), oFdSize, ACL_MEM_MALLOC_HUGE_FIRST));
@@ -268,26 +301,39 @@ static void Run(const Options &options)
268 ACL_CHECK(aclrtGetHardwareSyncAddr(reinterpret_cast<void**>(&hardwareSyncAddr)));301 ACL_CHECK(aclrtGetHardwareSyncAddr(reinterpret_cast<void**>(&hardwareSyncAddr)));
269 302 
270 // use Tp1Spec kernel to get better performance when numHeads = 128303 // use Tp1Spec kernel to get better performance when numHeads = 128
271- if (tilingKey == 0) {304+ switch (tilingKey) {
272- MLAFp16<<<blockDim, nullptr, stream>>>(305+ case 0:
273- hardwareSyncAddr, qDevice, qRopeDevice, kDevice, kRopeDevice, blockTableDevice, oDevice, sDevice, pDevice,306+ MLA<half><<<blockDim, nullptr, stream>>>(hardwareSyncAddr, qDevice, qRopeDevice, kDevice, kRopeDevice,
274- oTmpDevice, globaloDevice, oCoreTmpDevice, lDevice, tilingDevice307+ blockTableDevice, oDevice, sDevice, pDevice, oTmpDevice,
275- );308+ globaloDevice, oCoreTmpDevice, lDevice, tilingDevice);
276- } else if (tilingKey == 1) {309+ break;
277- MLABf16<<<blockDim, nullptr, stream>>>(310+ case 1:
278- hardwareSyncAddr, qDevice, qRopeDevice, kDevice, kRopeDevice, blockTableDevice, oDevice, sDevice, pDevice,311+ MLA<bfloat16_t><<<blockDim, nullptr, stream>>>(hardwareSyncAddr, qDevice, qRopeDevice, kDevice, kRopeDevice,
279- oTmpDevice, globaloDevice, oCoreTmpDevice, lDevice, tilingDevice312+ blockTableDevice, oDevice, sDevice, pDevice, oTmpDevice,
280- );313+ globaloDevice, oCoreTmpDevice, lDevice, tilingDevice);
281- } else if (tilingKey == 4) {314+ break;
282- MLATp1SpecFp16<<<blockDim, nullptr, stream>>>(315+ case 4:
283- hardwareSyncAddr, qDevice, qRopeDevice, kDevice, kRopeDevice, blockTableDevice, oDevice, sDevice, pDevice,316+ AMLATp1Spec<half><<<blockDim, nullptr, stream>>>(hardwareSyncAddr, qDevice, qRopeDevice, kDevice, kRopeDevice,
284- oTmpDevice, globaloDevice, oCoreTmpDevice, lDevice, tilingDevice317+ blockTableDevice, oDevice, sDevice, pDevice, oTmpDevice,
285- );318+ globaloDevice, oCoreTmpDevice, lDevice, tilingDevice);
286- } else if (tilingKey == 5) {319+ break;
287- MLATp1SpecBf16<<<blockDim, nullptr, stream>>>(320+ case 5:
288- hardwareSyncAddr, qDevice, qRopeDevice, kDevice, kRopeDevice, blockTableDevice, oDevice, sDevice, pDevice,321+ AMLATp1Spec<bfloat16_t><<<blockDim, nullptr, stream>>>(hardwareSyncAddr, qDevice, qRopeDevice, kDevice, kRopeDevice,
289- oTmpDevice, globaloDevice, oCoreTmpDevice, lDevice, tilingDevice322+ blockTableDevice, oDevice, sDevice, pDevice, oTmpDevice,
290- );323+ globaloDevice, oCoreTmpDevice, lDevice, tilingDevice);
324+ break;
325+ case 7:
326+ MLATp1Spec<half><<<blockDim, nullptr, stream>>>(hardwareSyncAddr, qDevice, qRopeDevice, kDevice, kRopeDevice,
327+ blockTableDevice, oDevice, sDevice, pDevice, oTmpDevice,
328+ globaloDevice, oCoreTmpDevice, lDevice, tilingDevice);
329+ break;
330+ case 8:
331+ MLATp1Spec<bfloat16_t><<<blockDim, nullptr, stream>>>(hardwareSyncAddr, qDevice, qRopeDevice, kDevice, kRopeDevice,
332+ blockTableDevice, oDevice, sDevice, pDevice, oTmpDevice,
333+ globaloDevice, oCoreTmpDevice, lDevice, tilingDevice);
334+ break;
335+ default:
336+ break;
291 }337 }
292 ACL_CHECK(aclrtSynchronizeStream(stream));338 ACL_CHECK(aclrtSynchronizeStream(stream));
293 // Copy the result from device to host339 // Copy the result from device to host
@@ -304,13 +350,28 @@ static void Run(const Options &options)
304 const size_t goldenSize = qoSize * 2;350 const size_t goldenSize = qoSize * 2;
305 ReadFile(dataPath + "/golden.bin", goldenHost.data(), goldenSize);351 ReadFile(dataPath + "/golden.bin", goldenHost.data(), goldenSize);
306 352 
307- // Compare the result353+ // Compute the cpulow result
308- vector<uint64_t> errorIndices = (dataType == "half") ? golden::CompareData(oHostHalf, goldenHost, kvSeqlen)354+ // cpu_low.bin is written in fp16 (same format as result.bin), NOT fp32 like golden.bin.
309- : golden::CompareData(oHostBf16, goldenHost, kvSeqlen);355+ // Read it as fp16 first, then widen to float for comparison.
310- if (errorIndices.empty()) {356+ vector<float> cpulowHost(qoSize / sizeof(fp16_t));
357+ vector<fp16_t> cpulowFp16Tmp(qoSize / sizeof(fp16_t));
358+ const size_t cpulowSize = qoSize; // fp16 binary: qoSize bytes
359+ ReadFile(dataPath + "/cpu_low.bin", cpulowFp16Tmp.data(), cpulowSize);
360+ for (size_t i = 0; i < cpulowFp16Tmp.size(); ++i) {
361+ cpulowHost[i] = static_cast<float>(cpulowFp16Tmp[i]);
362+ }
363+ 
364+ // Compute error metrics
365+ auto errorMetrics = (dataType == "half")
366+ ? golden::ComputeErrorMetrics(oHostHalf, cpulowHost, goldenHost, 10.0, 2.0, 2.0)
367+ : golden::ComputeErrorMetrics(oHostBf16, cpulowHost, goldenHost, 10.0, 2.0, 2.0);
368+ if (errorMetrics.passed) {
311 cout << "Compare success." << endl;369 cout << "Compare success." << endl;
312 } else {370 } else {
313- cerr << "Compare failed. Error count: " << errorIndices.size() << endl;371+ cerr << "Error ratios exceed thresholds:" << endl;
372+ cerr << "MARE ratio: " << errorMetrics.mareRatio << " (threshold: 10)" << endl;
373+ cerr << "MERE ratio: " << errorMetrics.mereRatio << " (threshold: 2)" << endl;
374+ cerr << "RMSE ratio: " << errorMetrics.rmseRatio << " (threshold: 2)" << endl;
314 }375 }
315 376 
316 // Free host memory allocations.377 // Free host memory allocations.
Mexamples/19_mla/mla.md+4-3
@@ -12,13 +12,13 @@ CATLASS MLA是基于CATLASS Gemm API实现的亲和昇腾AtlasA2硬件的Flash-M
12Tiling计算的逻辑位于[mla.cpp](./mla.cpp)文件中，在调用算子前，需要准备好tiling计算所需的各项参数，赋值给MLAInfo结构体，并调用`GetMLATilingParam`函数。[mla.cpp](./mla.cpp)中提供了一个示例12Tiling计算的逻辑位于[mla.cpp](./mla.cpp)文件中，在调用算子前，需要准备好tiling计算所需的各项参数，赋值给MLAInfo结构体，并调用`GetMLATilingParam`函数。[mla.cpp](./mla.cpp)中提供了一个示例
13 13 
14```c++14```c++
15-// 准备Tiling计算所需的中间结构体以及Host侧空间15+// 准备Tiling计算所需的中间结构体以及HOST侧空间
16MLATiling::MLAInfo mlaInfo;16MLATiling::MLAInfo mlaInfo;
17...17...
18MLATiling::GetMLATilingParam(mlaInfo, blockDim, (uint32_t *)tilingHost);18MLATiling::GetMLATilingParam(mlaInfo, blockDim, (uint32_t *)tilingHost);
19```19```
20 20 
21-`GetMLATilingParam`函数中，调用了两个函数`GetMLATilingCommon`与`GetMLATilingSpec`，分别对应了通用场景下和特化场景下的分核逻辑21+`GetMLATilingParam`函数中，调用了两个函数`GetMLATilingCommon`与`GetMLATilingSpec`，分别对应了通用场景下/特化场景下的分核逻辑
22 22 
23## Kernel23## Kernel
24 24 
@@ -42,7 +42,7 @@ MLATiling::GetMLATilingParam(mlaInfo, blockDim, (uint32_t *)tilingHost);
42在本算子中，使用了Block和Tile层级组件来组装Kernel，具体步骤为：42在本算子中，使用了Block和Tile层级组件来组装Kernel，具体步骤为：
43 43 
441. 组装attention计算中的两个BlockMmad（QK,PV）以及三个BlockEpilogue（softmax, rescaleO, flashDecoding）。441. 组装attention计算中的两个BlockMmad（QK,PV）以及三个BlockEpilogue（softmax, rescaleO, flashDecoding）。
45-2. 将Block组合在一起构建成`MLAKernel`，并在Kernel类中完成对各个Block的循环调用。 45+2. 将Block组合在一起构建成`MLAKernel`，并在Kernel类中完成对各个Block的循环调用。
46 46 
47这一过程也体现在Kernel入口的代码中（以[mla_kernel.cpp](./mla_kernel.cpp)为例）：47这一过程也体现在Kernel入口的代码中（以[mla_kernel.cpp](./mla_kernel.cpp)为例）：
48 48 
@@ -73,6 +73,7 @@ using EpilogueMLARescaleO =
73 Epilogue::Block::BlockEpilogue<Epilogue::EpilogueAtlasA2MLARescaleO, OType, OUpdateType, OTmpType>;73 Epilogue::Block::BlockEpilogue<Epilogue::EpilogueAtlasA2MLARescaleO, OType, OUpdateType, OTmpType>;
74 74 
75// Epilogue Block模块, 实现Flash MLA中flash decoding75// Epilogue Block模块, 实现Flash MLA中flash decoding
76+using OType = Gemm::GemmType<ElementO, LayoutO>;
76using lType = Gemm::GemmType<ElementUpdate, LayoutUpdate>;77using lType = Gemm::GemmType<ElementUpdate, LayoutUpdate>;
77constexpr uint32_t ComputeEleNum = 6144;78constexpr uint32_t ComputeEleNum = 6144;
78using EpilogueMLAFDRescaleO =79using EpilogueMLAFDRescaleO =
Mexamples/19_mla/mla_kernel.cpp+500-422
Mexamples/19_mla/mla_tiling.cpp+341-100
@@ -21,7 +21,6 @@
21#include <vector>21#include <vector>
22 22 
23#include "catlass/detail/alignment.hpp"23#include "catlass/detail/alignment.hpp"
24- 
25using namespace std;24using namespace std;
26namespace MLATiling {25namespace MLATiling {
27using AddrOffsets = struct AddressOffsetInfo {26using AddrOffsets = struct AddressOffsetInfo {
@@ -32,14 +31,17 @@ using AddrOffsets = struct AddressOffsetInfo {
32 uint64_t addrLSeqOffset = 0;31 uint64_t addrLSeqOffset = 0;
33};32};
34 33 
35-inline uint32_t GetHigh32Bit(uint64_t v) {34+struct AddrOffsetPerBatch {
36- return static_cast<uint32_t>(v >> NUM32);35+ uint64_t qSeqOffset = 0;
37-}36+ uint64_t qSeqRopeOffset = 0;
38-inline uint32_t GetLow32Bit(uint64_t v) {37+ uint64_t maskBatchOffset = 0;
39- return static_cast<uint32_t>(v);38+};
40-}
41 39 
42-void GetAddrOffsetMLA(uint32_t *tilingHost, const AddrOffsets addrOffsets, const int32_t tilingOffset) {40+inline uint32_t GetHigh32Bit(uint64_t v) { return static_cast<uint32_t>(v >> NUM32); }
41+inline uint32_t GetLow32Bit(uint64_t v) { return static_cast<uint32_t>(v); }
42+ 
43+void GetAddrOffsetMLA(uint32_t *tilingHost, const AddrOffsets &addrOffsets, const int32_t tilingOffset)
44+{
43 // Calculate address offset45 // Calculate address offset
44 tilingHost[tilingOffset + NUM4] = GetHigh32Bit(addrOffsets.addrQSeqOffset);46 tilingHost[tilingOffset + NUM4] = GetHigh32Bit(addrOffsets.addrQSeqOffset);
45 tilingHost[tilingOffset + NUM5] = GetLow32Bit(addrOffsets.addrQSeqOffset);47 tilingHost[tilingOffset + NUM5] = GetLow32Bit(addrOffsets.addrQSeqOffset);
@@ -49,46 +51,58 @@ void GetAddrOffsetMLA(uint32_t *tilingHost, const AddrOffsets addrOffsets, const
49 tilingHost[tilingOffset + NUM9] = GetLow32Bit(addrOffsets.addrMaskBatchOffset);51 tilingHost[tilingOffset + NUM9] = GetLow32Bit(addrOffsets.addrMaskBatchOffset);
50}52}
51 53 
52-void GetMLATilingCommon(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *tilingHost) {54+void GetMLATilingCommon(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *tilingHost, const std::vector<uint32_t>& sortedIndices,
53- // Calculate the batch-related tiling parameters55+ const std::vector<AddrOffsetPerBatch>& addrOffsetsPerBatch)
56+{
54 int32_t maxKVSeqlen = 0;57 int32_t maxKVSeqlen = 0;
55 int32_t maxQSeqlen = 0;58 int32_t maxQSeqlen = 0;
56- AddrOffsets addrOffsets{};
57 for (int32_t seqIdx = 0; seqIdx < mlaInfo.batch; seqIdx++) {59 for (int32_t seqIdx = 0; seqIdx < mlaInfo.batch; seqIdx++) {
58- int32_t qSeqLen = *(mlaInfo.qSeqLen + seqIdx);60+ uint32_t sortSeqIdx = sortedIndices[seqIdx];
59- qSeqLen = (*(mlaInfo.kvSeqLen + seqIdx) == 0) ? 0 : qSeqLen;61+ int32_t qSeqLen = *(mlaInfo.qSeqLen + sortSeqIdx);
62+ int32_t kvSeqlen = *(mlaInfo.kvSeqLen + sortSeqIdx);
63+ 
64+ qSeqLen = (kvSeqlen == 0) ? 0 : qSeqLen;
60 maxQSeqlen = std::max(maxQSeqlen, qSeqLen);65 maxQSeqlen = std::max(maxQSeqlen, qSeqLen);
61- int32_t kvSeqlen = *(mlaInfo.kvSeqLen + seqIdx);
62 maxKVSeqlen = std::max(maxKVSeqlen, kvSeqlen);66 maxKVSeqlen = std::max(maxKVSeqlen, kvSeqlen);
67+ 
63 int32_t tilingOffset = TILING_HEAD_SIZE + TILING_PARA_SIZE * seqIdx;68 int32_t tilingOffset = TILING_HEAD_SIZE + TILING_PARA_SIZE * seqIdx;
64 tilingHost[tilingOffset] = static_cast<uint32_t>(qSeqLen);69 tilingHost[tilingOffset] = static_cast<uint32_t>(qSeqLen);
65 tilingHost[tilingOffset + NUM1] = static_cast<uint32_t>(kvSeqlen);70 tilingHost[tilingOffset + NUM1] = static_cast<uint32_t>(kvSeqlen);
71+ tilingHost[tilingOffset + NUM2] = static_cast<uint32_t>(sortSeqIdx);
66 tilingHost[tilingOffset + NUM3] = static_cast<uint32_t>(mlaInfo.blockSize);72 tilingHost[tilingOffset + NUM3] = static_cast<uint32_t>(mlaInfo.blockSize);
73+ 
74+ const auto &batchOffsets = addrOffsetsPerBatch[sortSeqIdx];
75+ AddrOffsets addrOffsets{};
76+ addrOffsets.addrQSeqOffset = batchOffsets.qSeqOffset;
77+ addrOffsets.addrQSeqRopeOffset = batchOffsets.qSeqRopeOffset;
78+ addrOffsets.addrMaskBatchOffset = batchOffsets.maskBatchOffset;
67 GetAddrOffsetMLA(tilingHost, addrOffsets, tilingOffset);79 GetAddrOffsetMLA(tilingHost, addrOffsets, tilingOffset);
68- uint64_t addressOffset = static_cast<uint64_t>(mlaInfo.numHeads * mlaInfo.embeddingSize * qSeqLen);
69- uint64_t addressMaskOffset = static_cast<uint64_t>(mlaInfo.maxKvSeqlen * qSeqLen);
70- uint64_t addressOffsetRope = static_cast<uint64_t>(mlaInfo.numHeads * mlaInfo.embeddingSizeRope * qSeqLen);
71- addrOffsets.addrQSeqOffset += addressOffset;
72- addrOffsets.addrQSeqRopeOffset += addressOffsetRope;
73- addrOffsets.addrMaskBatchOffset += addressMaskOffset;
74 }80 }
81+ 
75 tilingHost[TILING_MAX_KVSEQLEN] = maxKVSeqlen;82 tilingHost[TILING_MAX_KVSEQLEN] = maxKVSeqlen;
76 tilingHost[TILING_MAX_QSEQLEN] = maxQSeqlen;83 tilingHost[TILING_MAX_QSEQLEN] = maxQSeqlen;
77}84}
78 85 
79-void GetMLATilingSpec(const MLAInfo &mmInfo, uint32_t &blockDim, uint32_t *tilingHost) {86+void GetMLATilingSpec(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *tilingHost, const std::vector<uint32_t>& sortedIndices,
80- // Tp1 senario specialization87+ const std::vector<std::vector<uint32_t>>& realPreTaskNums)
88+{
89+ // TP1 scenario specialization.
81 // Treat every Q token with 128 heads as one process, regardless of the mtp depth90 // Treat every Q token with 128 heads as one process, regardless of the mtp depth
82 int32_t prevTaskNum = 0;91 int32_t prevTaskNum = 0;
83 int32_t maxKVSeqlen = 0;92 int32_t maxKVSeqlen = 0;
84- for (int32_t seqIdx = 0; seqIdx < mmInfo.batch; seqIdx++) {93+ 
85- int32_t qSeqLen = mmInfo.qSeqLen == nullptr ? 1 : *(mmInfo.qSeqLen + seqIdx);94+ for (int32_t i = 0; i < mlaInfo.batch; i++) {
86- int32_t kvSeqlen = *(mmInfo.kvSeqLen + seqIdx);95+ uint32_t sortSeqIdx = sortedIndices[i];
87- maxKVSeqlen = std::max(maxKVSeqlen, kvSeqlen);96+ int32_t qSeqLen = *(mlaInfo.qSeqLen + sortSeqIdx);
97+ int32_t kvSeqlen = *(mlaInfo.kvSeqLen + sortSeqIdx);
98+ const std::vector<uint32_t>& currentRealPreTaskNum = realPreTaskNums[sortSeqIdx];
99+ 
100+ maxKVSeqlen = (maxKVSeqlen > kvSeqlen) ? maxKVSeqlen : kvSeqlen;
101+ 
88 for (int32_t qSeq = 0; qSeq < qSeqLen; qSeq++) {102 for (int32_t qSeq = 0; qSeq < qSeqLen; qSeq++) {
89 int32_t tilingOffset = TILING_HEAD_SIZE + PARA_TILING_ELENUM_SPEC * prevTaskNum;103 int32_t tilingOffset = TILING_HEAD_SIZE + PARA_TILING_ELENUM_SPEC * prevTaskNum;
90- tilingHost[tilingOffset] = seqIdx;104+ tilingHost[tilingOffset] = sortSeqIdx;
91- tilingHost[tilingOffset + NUM1] = prevTaskNum;105+ tilingHost[tilingOffset + NUM1] = currentRealPreTaskNum[qSeq];
92 tilingHost[tilingOffset + NUM2] = kvSeqlen;106 tilingHost[tilingOffset + NUM2] = kvSeqlen;
93 prevTaskNum++;107 prevTaskNum++;
94 }108 }
@@ -96,9 +110,10 @@ void GetMLATilingSpec(const MLAInfo &mmInfo, uint32_t &blockDim, uint32_t *tilin
96 tilingHost[TILING_MAX_KVSEQLEN] = maxKVSeqlen;110 tilingHost[TILING_MAX_KVSEQLEN] = maxKVSeqlen;
97}111}
98 112 
99-int32_t GetQNBlockTile(const MLAInfo &mlaInfo, int32_t qSeqLen, uint32_t specStrategyFlag) {113+int32_t GetQNBlockTile(const MLAInfo &mlaInfo, int32_t qSeqLen, uint32_t specStrategyFlag)
114+{
100 int32_t tokenNum = qSeqLen;115 int32_t tokenNum = qSeqLen;
101- if (specStrategyFlag > 0) {116+ if (specStrategyFlag) {
102 tokenNum = NUM1;117 tokenNum = NUM1;
103 }118 }
104 int32_t tileListIdx = static_cast<int32_t>(std::ceil(std::log2(tokenNum)));119 int32_t tileListIdx = static_cast<int32_t>(std::ceil(std::log2(tokenNum)));
@@ -109,17 +124,13 @@ int32_t GetQNBlockTile(const MLAInfo &mlaInfo, int32_t qSeqLen, uint32_t specStr
109 return qNBlockTile;124 return qNBlockTile;
110}125}
111 126 
112-void GetTilingHead(127+void GetTilingHead(const MLAInfo &mlaInfo, uint32_t *tilingHost, const uint32_t *torPtr, int32_t maxQseqlen,
113- const MLAInfo &mlaInfo,128+ uint32_t specStrategyFlag)
114- uint32_t *tilingHost,129+{
115- const uint32_t *torPtr,
116- int32_t maxQseqlen,
117- uint32_t specStrategyFlag
118-) {
119 // Calculating tiling parameters130 // Calculating tiling parameters
120 tilingHost[TILING_BATCH] = static_cast<uint32_t>(mlaInfo.batch);131 tilingHost[TILING_BATCH] = static_cast<uint32_t>(mlaInfo.batch);
121 tilingHost[TILING_HEADSIZE] = static_cast<uint32_t>(TILING_HEAD_SIZE);132 tilingHost[TILING_HEADSIZE] = static_cast<uint32_t>(TILING_HEAD_SIZE);
122- if (specStrategyFlag > 0) {133+ if (specStrategyFlag) {
123 tilingHost[TILING_PARASIZE] = static_cast<uint32_t>(PARA_TILING_ELENUM_SPEC);134 tilingHost[TILING_PARASIZE] = static_cast<uint32_t>(PARA_TILING_ELENUM_SPEC);
124 } else {135 } else {
125 tilingHost[TILING_PARASIZE] = static_cast<uint32_t>(TILING_PARA_SIZE);136 tilingHost[TILING_PARASIZE] = static_cast<uint32_t>(TILING_PARA_SIZE);
@@ -141,115 +152,303 @@ void GetTilingHead(
141 tilingHost[TILING_TOTAL_QTOKENS] = static_cast<uint32_t>(mlaInfo.numTokens);152 tilingHost[TILING_TOTAL_QTOKENS] = static_cast<uint32_t>(mlaInfo.numTokens);
142}153}
143 154 
144-uint32_t GetKVSplitParam(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *tilingHost) {155+uint32_t GetKVSplitParam(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *tilingHost, const std::vector<uint32_t>& sortedIndices)
145- // Calculate the tiling parameters related to flash decoding156+{
146- bool isKVSplit = (tilingHost[TILING_MAX_KVSEQLEN] >= blockDim * KV_SEQLEN_SLICE * NUM2)157+ // Only split KV when the longest sequence is long enough to be worth distributing across
147- && (tilingHost[TILING_BATCH] <= blockDim * SPLITKV_RATION && tilingHost[TILING_MAX_QSEQLEN] == 1);158+ // cores (>= blockDim * KV_SEQLEN_SLICE * 2). For short sequences, splitting adds extra
148- if (tilingHost[TILING_NUMHEADS] == NUM128 || !isKVSplit) {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);
164+ 
165+ if (!isKVSplit) {
149 tilingHost[TILING_KVCORENUM] = 1;166 tilingHost[TILING_KVCORENUM] = 1;
150- tilingHost[TILING_KVSPLIT] = tilingHost[TILING_MAX_KVSEQLEN];167+ tilingHost[TILING_PROCESSNUM] = 0;
151- std::cout << "TILING_KVSPLIT = " << tilingHost[TILING_KVSPLIT] << std::endl;168+ for (int32_t seqIdx = 0; seqIdx < mlaInfo.batch; seqIdx++) {
152- std::cout << "TILING_KVCORENUM = " << tilingHost[TILING_KVCORENUM] << std::endl;169+ uint32_t sortSeqIdx = sortedIndices[seqIdx];
170+ uint32_t kvSeqLen = *(mlaInfo.kvSeqLen + sortSeqIdx);
171+ int32_t tilingOffset = seqIdx * TILING_PARA_SIZE + TILING_HEAD_SIZE;
172+ tilingHost[tilingOffset + NUM15] = kvSeqLen;
173+ tilingHost[tilingOffset + NUM16] = 1;
174+ }
153 return tilingHost[TILING_BATCH];175 return tilingHost[TILING_BATCH];
154 }176 }
155 177 
156- uint32_t decoderBatch = tilingHost[TILING_BATCH];178+ // T+1-style balanced KV split: distribute available cores among the (q==1) decode tasks
157- uint32_t process = std::lcm(decoderBatch, blockDim);179+ // by repeatedly splitting the currently heaviest task, with a minimum tokens-per-split floor.
158- uint32_t kvSplitCoreNum = process / decoderBatch;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);
159 188 
160- uint32_t kvSeqlenMaxAlign = RoundUp(tilingHost[TILING_MAX_KVSEQLEN], static_cast<uint32_t>(mlaInfo.blockSize));189+ std::vector<uint32_t> alignedKvLens(taskCount);
161- uint32_t kvSeqBlockNum = kvSeqlenMaxAlign / mlaInfo.blockSize;190+ for (uint32_t seqIdx = 0; seqIdx < taskCount; seqIdx++) {
162- uint32_t kvBlockPerCore = CeilDiv(kvSeqBlockNum, kvSplitCoreNum);191+ uint32_t sortSeqIdx = sortedIndices[seqIdx];
163- uint32_t kvSplitPerCore = kvBlockPerCore * mlaInfo.blockSize;192+ uint32_t kvSeqLen = *(mlaInfo.kvSeqLen + sortSeqIdx);
164- kvSplitCoreNum = CeilDiv(tilingHost[TILING_MAX_KVSEQLEN], kvSplitPerCore);193+ alignedKvLens[seqIdx] = RoundUp(kvSeqLen, static_cast<uint32_t>(mlaInfo.blockSize));
194+ }
165 195 
166- tilingHost[TILING_KVSPLIT] = kvSplitPerCore;196+ std::vector<uint32_t> batchSplitNum(taskCount, 1);
167- tilingHost[TILING_KVCORENUM] = kvSplitCoreNum;197+ uint32_t totalAllocated = taskCount;
168- std::cout << "TILING_KVSPLIT = " << tilingHost[TILING_KVSPLIT] << std::endl;198+ while (totalAllocated < blockDim) {
169- std::cout << "TILING_KVCORENUM = " << tilingHost[TILING_KVCORENUM] << std::endl;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 &&
205+ nextLoad >= MIN_TOKENS_PER_SPLIT && currentLoad > maxLoad) {
206+ maxLoad = currentLoad;
207+ maxLoadIdx = static_cast<int32_t>(i);
208+ }
209+ }
210+ if (maxLoadIdx == -1) {
211+ break;
212+ }
213+ batchSplitNum[maxLoadIdx]++;
214+ totalAllocated++;
215+ }
216+ 
217+ uint32_t MAX_KV_SPLIT_NUM = 1;
218+ for (uint32_t num : batchSplitNum) {
219+ if (num > MAX_KV_SPLIT_NUM) {
220+ MAX_KV_SPLIT_NUM = num;
221+ }
222+ }
223+ 
224+ // Build the cumulative-task prefix sum so the kernel can map process -> (task, kvIdx) tightly.
225+ uint32_t cuTaskVal = 0;
226+ int32_t cuTaskIdx = 0;
227+ tilingHost[CUTASK_START_OFFSET + cuTaskIdx++] = cuTaskVal;
228+ for (uint32_t seqIdx = 0; seqIdx < taskCount; seqIdx++) {
229+ int32_t tilingOffset = seqIdx * TILING_PARA_SIZE + TILING_HEAD_SIZE;
230+ uint32_t sortSeqIdx = sortedIndices[seqIdx];
231+ uint32_t kvSeqLen = *(mlaInfo.kvSeqLen + sortSeqIdx);
232+ 
233+ uint32_t kvSplitPerCore;
234+ uint32_t kvSplitCoreNum;
235+ if (kvSeqLen == 0) {
236+ kvSplitPerCore = static_cast<uint32_t>(mlaInfo.blockSize);
237+ kvSplitCoreNum = 1;
238+ } else {
239+ uint32_t allocSplit = batchSplitNum[seqIdx];
240+ uint32_t kvSeqAlign = RoundUp(kvSeqLen, static_cast<uint32_t>(mlaInfo.blockSize));
241+ uint32_t kvSeqBlockNum = kvSeqAlign / mlaInfo.blockSize;
242+ uint32_t kvBlockPerCore = CeilDiv(kvSeqBlockNum, allocSplit);
243+ kvSplitPerCore = kvBlockPerCore * mlaInfo.blockSize;
244+ kvSplitCoreNum = CeilDiv(kvSeqLen, kvSplitPerCore);
245+ }
246+ tilingHost[tilingOffset + NUM15] = kvSplitPerCore;
247+ tilingHost[tilingOffset + NUM16] = kvSplitCoreNum;
248+ cuTaskVal += kvSplitCoreNum;
249+ tilingHost[CUTASK_START_OFFSET + cuTaskIdx++] = cuTaskVal;
250+ }
251+ tilingHost[TILING_PROCESSNUM] = cuTaskVal;
252+ tilingHost[TILING_KVCORENUM] = MAX_KV_SPLIT_NUM;
253+ std::cout << "TILING_MAX_KVCORENUM = " << tilingHost[TILING_KVCORENUM] << std::endl;
254+ std::cout << "TILING_PROCESSNUM = " << tilingHost[TILING_PROCESSNUM] << std::endl;
170 255 
171 // Set lOffsetInfo and OfdOffsetInfo256 // Set lOffsetInfo and OfdOffsetInfo
172 AddrOffsets addrOffsets;257 AddrOffsets addrOffsets;
173 for (int32_t seqIdx = 0; seqIdx < mlaInfo.batch; seqIdx++) {258 for (int32_t seqIdx = 0; seqIdx < mlaInfo.batch; seqIdx++) {
174- int32_t qSeqlen = 1;259+ uint32_t sortSeqIdx = sortedIndices[seqIdx];
175- qSeqlen = (*(mlaInfo.kvSeqLen + seqIdx) == 0) ? 0 : qSeqlen;260+ uint32_t qSeqlen = *(mlaInfo.qSeqLen + sortSeqIdx);
261+ uint32_t kvSeqLen = *(mlaInfo.kvSeqLen + sortSeqIdx);
262+ qSeqlen = (kvSeqLen == 0) ? 0 : qSeqlen;
176 int32_t tilingOffset = seqIdx * TILING_PARA_SIZE + TILING_HEAD_SIZE;263 int32_t tilingOffset = seqIdx * TILING_PARA_SIZE + TILING_HEAD_SIZE;
177 tilingHost[tilingOffset + NUM11] = GetHigh32Bit(addrOffsets.addrLSeqOffset);264 tilingHost[tilingOffset + NUM11] = GetHigh32Bit(addrOffsets.addrLSeqOffset);
178 tilingHost[tilingOffset + NUM12] = GetLow32Bit(addrOffsets.addrLSeqOffset);265 tilingHost[tilingOffset + NUM12] = GetLow32Bit(addrOffsets.addrLSeqOffset);
179 tilingHost[tilingOffset + NUM13] = GetHigh32Bit(addrOffsets.addrOFdSeqOffset);266 tilingHost[tilingOffset + NUM13] = GetHigh32Bit(addrOffsets.addrOFdSeqOffset);
180 tilingHost[tilingOffset + NUM14] = GetLow32Bit(addrOffsets.addrOFdSeqOffset);267 tilingHost[tilingOffset + NUM14] = GetLow32Bit(addrOffsets.addrOFdSeqOffset);
181- addrOffsets.addrLSeqOffset += static_cast<uint64_t>(mlaInfo.numHeads * qSeqlen * kvSplitCoreNum);268+ addrOffsets.addrLSeqOffset += static_cast<uint64_t>(mlaInfo.numHeads * qSeqlen * MAX_KV_SPLIT_NUM);
182 addrOffsets.addrOFdSeqOffset += static_cast<uint64_t>(mlaInfo.numHeads * qSeqlen * mlaInfo.embeddingSize);269 addrOffsets.addrOFdSeqOffset += static_cast<uint64_t>(mlaInfo.numHeads * qSeqlen * mlaInfo.embeddingSize);
183 }270 }
184 271 
185- return decoderBatch * kvSplitCoreNum;272+ return cuTaskVal;
186}273}
187 274 
188-uint32_t GetKVSplitParamSpec(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *tilingHost) {275+uint32_t GetKVSplitParamSpec(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *tilingHost, const std::vector<uint32_t>& sortedIndices)
189- // Tp1 senario specialization276+{
277+ // TP1 scenario specialization.
190 // Calculate the tiling parameters related to flash decoding278 // Calculate the tiling parameters related to flash decoding
191 uint32_t totalTaskNumSpec = tilingHost[TILING_TOTAL_QTOKENS];279 uint32_t totalTaskNumSpec = tilingHost[TILING_TOTAL_QTOKENS];
192- 
193 uint32_t formerTaskNum = totalTaskNumSpec;280 uint32_t formerTaskNum = totalTaskNumSpec;
194 uint32_t tailTaskNum = 0;281 uint32_t tailTaskNum = 0;
195- 282+ if (mlaInfo.numTokens % NUM20 <= NUM10 && mlaInfo.batch <= 40) {
196- uint32_t processLoop = totalTaskNumSpec / blockDim;283+ uint32_t processLoop = totalTaskNumSpec / blockDim;
197- formerTaskNum = processLoop * blockDim;284+ formerTaskNum = processLoop * blockDim;
198- tailTaskNum = totalTaskNumSpec - formerTaskNum;285+ tailTaskNum = totalTaskNumSpec - formerTaskNum;
199- 
200- if (tailTaskNum >= blockDim * SPLITKV_RATION) {
201- formerTaskNum = totalTaskNumSpec;
202- tailTaskNum = 0;
203 }286 }
204 287 
205 tilingHost[TILING_FORMERTASKNUM] = formerTaskNum;288 tilingHost[TILING_FORMERTASKNUM] = formerTaskNum;
206 tilingHost[TILING_TAILTASKNUM] = tailTaskNum;289 tilingHost[TILING_TAILTASKNUM] = tailTaskNum;
207 std::cout << "TILING_FORMERTASKNUM = " << tilingHost[TILING_FORMERTASKNUM] << std::endl;290 std::cout << "TILING_FORMERTASKNUM = " << tilingHost[TILING_FORMERTASKNUM] << std::endl;
208 std::cout << "TILING_TAILTASKNUM = " << tilingHost[TILING_TAILTASKNUM] << std::endl;291 std::cout << "TILING_TAILTASKNUM = " << tilingHost[TILING_TAILTASKNUM] << std::endl;
209- 
210 if (tailTaskNum == 0) {292 if (tailTaskNum == 0) {
211 tilingHost[TILING_KVCORENUM] = 1;293 tilingHost[TILING_KVCORENUM] = 1;
212 tilingHost[TILING_KVSPLIT] = tilingHost[TILING_MAX_KVSEQLEN];294 tilingHost[TILING_KVSPLIT] = tilingHost[TILING_MAX_KVSEQLEN];
295+ tilingHost[TILING_PROCESSNUM] = 0;
213 std::cout << "TILING_KVSPLIT = " << tilingHost[TILING_KVSPLIT] << std::endl;296 std::cout << "TILING_KVSPLIT = " << tilingHost[TILING_KVSPLIT] << std::endl;
214 std::cout << "TILING_KVCORENUM = " << tilingHost[TILING_KVCORENUM] << std::endl;297 std::cout << "TILING_KVCORENUM = " << tilingHost[TILING_KVCORENUM] << std::endl;
215 return blockDim;298 return blockDim;
216 }299 }
217 300 
218- uint32_t process = std::lcm(tailTaskNum, blockDim);301+ const uint32_t MIN_TOKENS_PER_SPLIT = static_cast<uint32_t>(mlaInfo.blockSize * NUM4);
219- uint32_t kvSplitCoreNum = process / tailTaskNum;
220 302 
221- uint32_t kvSeqlenMaxAlign = RoundUp(tilingHost[TILING_MAX_KVSEQLEN], static_cast<uint32_t>(mlaInfo.blockSize));303+ std::vector<uint32_t> alignedKvLens;
222- uint32_t kvSeqBlockNum = kvSeqlenMaxAlign / mlaInfo.blockSize;304+ int32_t prevTaskNumTemp = 0;
223- uint32_t kvBlockPerCore = CeilDiv(kvSeqBlockNum, kvSplitCoreNum);
224- uint32_t kvSplitPerCore = kvBlockPerCore * mlaInfo.blockSize;
225- kvSplitCoreNum = CeilDiv(tilingHost[TILING_MAX_KVSEQLEN], kvSplitPerCore);
226- 
227- tilingHost[TILING_KVSPLIT] = kvSplitPerCore;
228- tilingHost[TILING_KVCORENUM] = kvSplitCoreNum;
229- std::cout << "TILING_KVSPLIT = " << tilingHost[TILING_KVSPLIT] << std::endl;
230- std::cout << "TILING_KVCORENUM = " << tilingHost[TILING_KVCORENUM] << std::endl;
231- 
232- // Set lOffsetInfo and OfdOffsetInfo
233- AddrOffsets addrOffsets;
234- int32_t prevTaskNum = 0;
235 for (int32_t seqIdx = 0; seqIdx < mlaInfo.batch; seqIdx++) {305 for (int32_t seqIdx = 0; seqIdx < mlaInfo.batch; seqIdx++) {
236- int32_t qSeqLen = mlaInfo.qSeqLen == nullptr ? 1 : *(mlaInfo.qSeqLen + seqIdx);306+ uint32_t sortSeqIdx = sortedIndices[seqIdx];
307+ uint32_t qSeqLen = *(mlaInfo.qSeqLen + sortSeqIdx);
308+ if (prevTaskNumTemp + qSeqLen < formerTaskNum) {
309+ prevTaskNumTemp += qSeqLen;
310+ continue;
311+ }
312+ uint32_t kvSeqLen = *(mlaInfo.kvSeqLen + sortSeqIdx);
313+ alignedKvLens.push_back(RoundUp(kvSeqLen, static_cast<uint32_t>(mlaInfo.blockSize)));
314+ }
315+ 
316+ uint32_t tailSeqCount = alignedKvLens.size();
317+ std::vector<uint32_t> batchSplitNum(tailSeqCount, 1);
318+ uint32_t totalAllocated = tailTaskNum;
319+ while (totalAllocated < blockDim) {
320+ int32_t maxLoadIdx = -1;
321+ uint32_t maxLoad = 0;
322+ 
323+ for (int32_t i = 0; i < tailSeqCount; i++) {
324+ uint32_t currentLoad = alignedKvLens[i] / batchSplitNum[i];
325+ uint32_t nextLoad = alignedKvLens[i] / (batchSplitNum[i] + 1);
326+ if (nextLoad >= MIN_TOKENS_PER_SPLIT && currentLoad > maxLoad) {
327+ maxLoad = currentLoad;
328+ maxLoadIdx = i;
329+ }
330+ }
331+ 
332+ if (maxLoadIdx != -1) {
333+ batchSplitNum[maxLoadIdx]++;
334+ } else if (tailSeqCount > 0) {
335+ int32_t longestIdx = 0;
336+ for (int32_t i = 1; i < tailSeqCount; i++) {
337+ if (alignedKvLens[i] > alignedKvLens[longestIdx]) {
338+ longestIdx = i;
339+ }
340+ }
341+ batchSplitNum[longestIdx]++;
342+ }
343+ totalAllocated++;
344+ }
345+ uint32_t MAX_KV_SPLIT_NUM = 1;
346+ for (uint32_t num : batchSplitNum) {
347+ if (num > MAX_KV_SPLIT_NUM) {
348+ MAX_KV_SPLIT_NUM = num;
349+ }
350+ }
351+ uint32_t cuTaskVal = 0;
352+ int32_t cuTaskIdx = 0;
353+ tilingHost[CUTASK_START_OFFSET + cuTaskIdx++] = cuTaskVal;
354+ 
355+ int32_t prevTaskNum = 0;
356+ int32_t tailSeqIdx = 0;
357+ for (int32_t seqIdx = 0; seqIdx < mlaInfo.batch; seqIdx++) {
358+ uint32_t sortSeqIdx = sortedIndices[seqIdx];
359+ uint32_t qSeqLen = *(mlaInfo.qSeqLen + sortSeqIdx);
360+ uint32_t kvSeqLen = *(mlaInfo.kvSeqLen + sortSeqIdx);
361+ 
362+ if (prevTaskNum + qSeqLen < formerTaskNum) {
363+ prevTaskNum += qSeqLen;
364+ continue;
365+ }
366+ 
367+ uint32_t allocSplit = 1;
368+ if (tailSeqIdx < tailSeqCount) {
369+ allocSplit = batchSplitNum[tailSeqIdx++];
370+ }
371+ 
372+ uint32_t kvSeqAlign = RoundUp(kvSeqLen, static_cast<uint32_t>(mlaInfo.blockSize));
373+ uint32_t kvSeqBlockNum = kvSeqAlign / mlaInfo.blockSize;
374+ uint32_t kvBlockPerCore = CeilDiv(kvSeqBlockNum, allocSplit);
375+ uint32_t kvSplitPerCore = kvBlockPerCore * mlaInfo.blockSize;
376+ uint32_t kvSplitCoreNum = CeilDiv(kvSeqLen, kvSplitPerCore);
377+ 
378+ for (int32_t qSeq = 0; qSeq < qSeqLen; qSeq++) {
379+ if (prevTaskNum >= formerTaskNum) {
380+ cuTaskVal += kvSplitCoreNum;
381+ tilingHost[CUTASK_START_OFFSET + cuTaskIdx++] = cuTaskVal;
382+ int32_t tilingOffset = TILING_HEAD_SIZE + PARA_TILING_ELENUM_SPEC * prevTaskNum;
383+ tilingHost[tilingOffset + NUM15] = kvSplitPerCore;
384+ tilingHost[tilingOffset + NUM16] = kvSplitCoreNum;
385+ }
386+ prevTaskNum++;
387+ }
388+ }
389+ tilingHost[TILING_PROCESSNUM] = cuTaskVal;
390+ tilingHost[CUTASK_START_OFFSET + cuTaskIdx++] = TILING_HEAD_SIZE;
391+ tilingHost[TILING_KVCORENUM] = MAX_KV_SPLIT_NUM;
392+ std::cout << "TILING_MAX_KVCORENUM = " << tilingHost[TILING_KVCORENUM] << std::endl;
393+ 
394+ AddrOffsets addrOffsets{};
395+ prevTaskNum = 0;
396+ 
397+ for (int32_t seqIdx = 0; seqIdx < mlaInfo.batch; seqIdx++) {
398+ uint32_t sortSeqIdx = sortedIndices[seqIdx];
399+ uint32_t qSeqLen = *(mlaInfo.qSeqLen + sortSeqIdx);
237 for (int32_t qSeq = 0; qSeq < qSeqLen; qSeq++) {400 for (int32_t qSeq = 0; qSeq < qSeqLen; qSeq++) {
238 int32_t tilingOffset = TILING_HEAD_SIZE + PARA_TILING_ELENUM_SPEC * prevTaskNum;401 int32_t tilingOffset = TILING_HEAD_SIZE + PARA_TILING_ELENUM_SPEC * prevTaskNum;
239 tilingHost[tilingOffset + NUM11] = GetHigh32Bit(addrOffsets.addrLSeqOffset);402 tilingHost[tilingOffset + NUM11] = GetHigh32Bit(addrOffsets.addrLSeqOffset);
240 tilingHost[tilingOffset + NUM12] = GetLow32Bit(addrOffsets.addrLSeqOffset);403 tilingHost[tilingOffset + NUM12] = GetLow32Bit(addrOffsets.addrLSeqOffset);
241 tilingHost[tilingOffset + NUM13] = GetHigh32Bit(addrOffsets.addrOFdSeqOffset);404 tilingHost[tilingOffset + NUM13] = GetHigh32Bit(addrOffsets.addrOFdSeqOffset);
242 tilingHost[tilingOffset + NUM14] = GetLow32Bit(addrOffsets.addrOFdSeqOffset);405 tilingHost[tilingOffset + NUM14] = GetLow32Bit(addrOffsets.addrOFdSeqOffset);
243- addrOffsets.addrLSeqOffset += static_cast<uint64_t>(mlaInfo.numHeads * kvSplitCoreNum);406+ addrOffsets.addrLSeqOffset += static_cast<uint64_t>(mlaInfo.numHeads * MAX_KV_SPLIT_NUM);
244 addrOffsets.addrOFdSeqOffset += static_cast<uint64_t>(mlaInfo.numHeads * mlaInfo.embeddingSize);407 addrOffsets.addrOFdSeqOffset += static_cast<uint64_t>(mlaInfo.numHeads * mlaInfo.embeddingSize);
245 prevTaskNum++;408 prevTaskNum++;
246 }409 }
247 }410 }
248 411 
249- return tailTaskNum * kvSplitCoreNum;412+ return tailTaskNum * MAX_KV_SPLIT_NUM;
250}413}
251 414 
252-int32_t GetMLATilingParam(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *tilingHost) {415+void swapIndices(std::vector<uint32_t>& indices, int i, int j) {
416+ uint32_t temp = indices[i];
417+ indices[i] = indices[j];
418+ indices[j] = temp;
419+}
420+ 
421+int partitionIndices(std::vector<uint32_t>& indices, const int32_t* kvSeqLen, const int32_t* qSeqLen, int low, int high) {
422+ uint32_t pivotIndex = indices[high];
423+ int32_t pivotKv = kvSeqLen[pivotIndex];
424+ int32_t pivotQ = qSeqLen[pivotIndex];
425+ int i = (low - 1);
426+ 
427+ for (int j = low; j <= high - 1; j++) {
428+ uint32_t currIndex = indices[j];
429+ int32_t currKv = kvSeqLen[currIndex];
430+ int32_t currQ = qSeqLen[currIndex];
431+ 
432+ bool needSwap = (currKv > pivotKv) || (currKv == pivotKv && currQ > pivotQ);
433+ if (needSwap) {
434+ i++;
435+ swapIndices(indices, i, j);
436+ }
437+ }
438+ swapIndices(indices, i + 1, high);
439+ return (i + 1);
440+}
441+ 
442+void quickSortIndices(std::vector<uint32_t>& indices, const int32_t* kvSeqLen, const int32_t* qSeqLen, int low, int high) {
443+ if (low < high) {
444+ int pi = partitionIndices(indices, kvSeqLen, qSeqLen, low, high);
445+ quickSortIndices(indices, kvSeqLen, qSeqLen, low, pi - 1);
446+ quickSortIndices(indices, kvSeqLen, qSeqLen, pi + 1, high);
447+ }
448+}
449+ 
450+int32_t GetMLATilingParam(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *tilingHost)
451+{
253 if (tilingHost == nullptr || mlaInfo.qSeqLen == nullptr || mlaInfo.kvSeqLen == nullptr) {452 if (tilingHost == nullptr || mlaInfo.qSeqLen == nullptr || mlaInfo.kvSeqLen == nullptr) {
254 cerr << "[ERROR] pointer tilingHost or seq is nullptr." << endl;453 cerr << "[ERROR] pointer tilingHost or seq is nullptr." << endl;
255 return -1;454 return -1;
@@ -258,6 +457,47 @@ int32_t GetMLATilingParam(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *
258 cerr << "[ERROR] blockSize != 128 is not supported." << endl;457 cerr << "[ERROR] blockSize != 128 is not supported." << endl;
259 return -1;458 return -1;
260 }459 }
460+ 
461+ std::vector<AddrOffsetPerBatch> addrOffsetsPerBatch(mlaInfo.batch);
462+ uint64_t currentQSeqOffset = 0;
463+ uint64_t currentQSeqRopeOffset = 0;
464+ uint64_t currentMaskBatchOffset = 0;
465+ 
466+ for (int32_t i = 0; i < mlaInfo.batch; ++i) {
467+ int32_t qSeqLen = *(mlaInfo.qSeqLen + i);
468+ 
469+ uint64_t deltaQSeq = static_cast<uint64_t>(mlaInfo.numHeads * mlaInfo.embeddingSize * qSeqLen);
470+ uint64_t deltaQSeqRope = static_cast<uint64_t>(mlaInfo.numHeads * mlaInfo.embeddingSizeRope * qSeqLen);
471+ uint64_t deltaMaskBatch = static_cast<uint64_t>(mlaInfo.maxKvSeqlen * qSeqLen);
472+ 
473+ addrOffsetsPerBatch[i].qSeqOffset = currentQSeqOffset;
474+ addrOffsetsPerBatch[i].qSeqRopeOffset = currentQSeqRopeOffset;
475+ addrOffsetsPerBatch[i].maskBatchOffset = currentMaskBatchOffset;
476+ 
477+ currentQSeqOffset += deltaQSeq;
478+ currentQSeqRopeOffset += deltaQSeqRope;
479+ currentMaskBatchOffset += deltaMaskBatch;
480+ }
481+ 
482+ std::vector<uint32_t> sortedIndices(mlaInfo.batch);
483+ for (uint32_t i = 0; i < mlaInfo.batch; ++i) {
484+ sortedIndices[i] = i;
485+ }
486+ 
487+ std::vector<std::vector<uint32_t>> realPreTaskNums(mlaInfo.batch, std::vector<uint32_t>(NUM512));
488+ int32_t preTaskNum = 0;
489+ for (int32_t i = 0; i < mlaInfo.batch; ++i) {
490+ uint32_t qSeqlen = static_cast<uint32_t>(*(mlaInfo.qSeqLen + i));
491+ for (int32_t token = 0; token < qSeqlen; token++) {
492+ realPreTaskNums[i][token] = preTaskNum;
493+ preTaskNum++;
494+ }
495+ }
496+ 
497+ if (mlaInfo.numTokens > 0) {
498+ quickSortIndices(sortedIndices, mlaInfo.kvSeqLen, mlaInfo.qSeqLen, 0, mlaInfo.batch - 1);
499+ }
500+ 
261 int32_t maxQseqlen = 0;501 int32_t maxQseqlen = 0;
262 int32_t totalKvNumtokens = 0;502 int32_t totalKvNumtokens = 0;
263 for (int32_t seqIdx = 0; seqIdx < mlaInfo.batch; seqIdx++) {503 for (int32_t seqIdx = 0; seqIdx < mlaInfo.batch; seqIdx++) {
@@ -277,16 +517,17 @@ int32_t GetMLATilingParam(const MLAInfo &mlaInfo, uint32_t &blockDim, uint32_t *
277 float tor = static_cast<float>(1.0 / sqrt(1.0 * (mlaInfo.embeddingSize + mlaInfo.embeddingSizeRope)));517 float tor = static_cast<float>(1.0 / sqrt(1.0 * (mlaInfo.embeddingSize + mlaInfo.embeddingSizeRope)));
278 uint32_t *torPtr = reinterpret_cast<uint32_t *>(&tor);518 uint32_t *torPtr = reinterpret_cast<uint32_t *>(&tor);
279 uint32_t specStrategyFlag = (mlaInfo.numHeads == NUM128) ? 1 : 0;519 uint32_t specStrategyFlag = (mlaInfo.numHeads == NUM128) ? 1 : 0;
280- if (specStrategyFlag > 0) {520+ 
281- GetMLATilingSpec(mlaInfo, blockDim, tilingHost);521+ if (specStrategyFlag) {
522+ GetMLATilingSpec(mlaInfo, blockDim, tilingHost, sortedIndices, realPreTaskNums);
282 } else {523 } else {
283- GetMLATilingCommon(mlaInfo, blockDim, tilingHost);524+ GetMLATilingCommon(mlaInfo, blockDim, tilingHost, sortedIndices, addrOffsetsPerBatch);
284 }525 }
285 GetTilingHead(mlaInfo, tilingHost, torPtr, maxQseqlen, specStrategyFlag);526 GetTilingHead(mlaInfo, tilingHost, torPtr, maxQseqlen, specStrategyFlag);
286- if (specStrategyFlag > 0) {527+ if (specStrategyFlag) {
287- GetKVSplitParamSpec(mlaInfo, blockDim, tilingHost);528+ GetKVSplitParamSpec(mlaInfo, blockDim, tilingHost, sortedIndices);
288 } else {529 } else {
289- GetKVSplitParam(mlaInfo, blockDim, tilingHost);530+ GetKVSplitParam(mlaInfo, blockDim, tilingHost, sortedIndices);
290 }531 }
291 return 0;532 return 0;
292}533}
Mexamples/19_mla/mla_tiling.h+5-1
@@ -37,7 +37,11 @@ const int32_t TILING_TOTAL_QTOKENS = 18;
37const int32_t TILING_FORMERTASKNUM = 19;37const int32_t TILING_FORMERTASKNUM = 19;
38const int32_t TILING_TAILTASKNUM = 20;38const int32_t TILING_TAILTASKNUM = 20;
39 39 
40-const int32_t TILING_HEAD_SIZE = 24;40+const int32_t TILING_PROCESSNUM = 21;
41+ 
42+const int32_t TILING_HEAD_SIZE = 400;
43+const int32_t CUTASK_START_OFFSET = 25;
44+const int32_t CUTASK_MAX_LENGTH = 130;
41const int32_t TILING_PARA_SIZE = 17;45const int32_t TILING_PARA_SIZE = 17;
42 46 
43const int32_t PARA_TILING_ELENUM_SPEC = 17;47const int32_t PARA_TILING_ELENUM_SPEC = 17;
Minclude/catlass/epilogue/block/block_epilogue.hpp+2-0
@@ -38,6 +38,8 @@ class BlockEpilogue {
38#include "catlass/epilogue/block/block_epilogue_gemv.hpp"38#include "catlass/epilogue/block/block_epilogue_gemv.hpp"
39#include "catlass/epilogue/block/block_epilogue_mla_tp1_softmax.hpp"39#include "catlass/epilogue/block/block_epilogue_mla_tp1_softmax.hpp"
40#include "catlass/epilogue/block/block_epilogue_mla_tp1_rescale_o.hpp"40#include "catlass/epilogue/block/block_epilogue_mla_tp1_rescale_o.hpp"
41+#include "catlass/epilogue/block/block_epilogue_amla_tp1_softmax.hpp"
42+#include "catlass/epilogue/block/block_epilogue_amla_tp1_rescale_o.hpp"
41#include "catlass/epilogue/block/block_epilogue_online_softmax_no_mask.hpp"43#include "catlass/epilogue/block/block_epilogue_online_softmax_no_mask.hpp"
42#include "catlass/epilogue/block/block_epilogue_rescale_o_no_split_row.hpp"44#include "catlass/epilogue/block/block_epilogue_rescale_o_no_split_row.hpp"
43#include "catlass/epilogue/block/block_epilogue_w4a4_per_token_per_channel_dequant.hpp"45#include "catlass/epilogue/block/block_epilogue_w4a4_per_token_per_channel_dequant.hpp"
Ainclude/catlass/epilogue/block/block_epilogue_amla_tp1_rescale_o.hpp+267-0
@@ -0,0 +1,267 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_AMLA_TP1_RESCALE_O_HPP
12+#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_AMLA_TP1_RESCALE_O_HPP
13+ 
14+#include "catlass/catlass.hpp"
15+#include "catlass/arch/resource.hpp"
16+#include "catlass/epilogue/dispatch_policy.hpp"
17+#include "catlass/epilogue/tile/tile_copy.hpp"
18+#include "catlass/gemm_coord.hpp"
19+#include "catlass/matrix_coord.hpp"
20+ 
21+namespace Catlass::Epilogue::Block {
22+ 
23+template <class OutputType_, class UpdateType_, class InputType_>
24+class BlockEpilogue<EpilogueAtlasA2AMLATP1RescaleO, OutputType_, UpdateType_, InputType_> {
25+public:
26+ // Type aliases
27+ using DispatchPolicy = EpilogueAtlasA2AMLATP1RescaleO;
28+ using ArchTag = typename DispatchPolicy::ArchTag;
29+ 
30+ using ElementOutput = typename OutputType_::Element;
31+ using ElementUpdate = typename UpdateType_::Element;
32+ using ElementInput = typename InputType_::Element;
33+ 
34+ using LayoutOutput = typename OutputType_::Layout;
35+ using LayoutUpdate = typename UpdateType_::Layout;
36+ using LayoutInput = typename InputType_::Layout;
37+ 
38+ static constexpr uint32_t HALF_ELENUM_PER_BLK = 16;
39+ static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128;
40+ static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64;
41+ static constexpr uint32_t HALF_ELENUM_PER_LINE = 256;
42+ static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128;
43+ static constexpr uint32_t MULTIPLIER = 2;
44+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
45+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
46+ static constexpr uint32_t UB_UINT8_LINE_SIZE = 512;
47+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE_MLA = 16384;
48+ static constexpr uint32_t ROW_WISE_CYCLE_TILE = 8;
49+ static constexpr uint32_t HALF_DM_UB_SIZE = 64;
50+ static constexpr uint32_t HALF_LL_UB_SIZE = 256;
51+ static constexpr uint32_t VECTOR_SIZE = 128;
52+ static constexpr uint32_t NUM4 = 4;
53+ 
54+ CATLASS_DEVICE
55+ BlockEpilogue(Arch::Resource<ArchTag> &resource, uint32_t kvSplitCoreNum_ = 1)
56+ {
57+ // Allocate UB space
58+ constexpr uint32_t LO_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE_MLA;
59+ constexpr uint32_t DM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 2 * UB_UINT8_LINE_SIZE;
60+ constexpr uint32_t GL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 7 * UB_UINT8_LINE_SIZE;
61+ constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE_MLA;
62+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA;
63+ constexpr uint32_t HM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE;
64+ constexpr uint32_t GM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 5 * UB_UINT8_LINE_SIZE;
65+ 
66+ kvSplitCoreNum = kvSplitCoreNum_;
67+ loUbTensor = resource.ubBuf.template GetBufferByByte<float>(LO_UB_TENSOR_OFFSET);
68+ dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
69+ glUbTensor[0] = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
70+ glUbTensor[1] = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET + UB_UINT8_LINE_SIZE);
71+ tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
72+ goUbTensor16 = resource.ubBuf.template GetBufferByByte<ElementOutput>(GO_UB_TENSOR_OFFSET);
73+ goUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(GO_UB_TENSOR_OFFSET);
74+ hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);
75+ gmUbTensor[0] = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
76+ gmUbTensor[1] = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET + UB_UINT8_LINE_SIZE);
77+ }
78+ 
79+ CATLASS_DEVICE
80+ ~BlockEpilogue() {}
81+ 
82+ CATLASS_DEVICE
83+ void SetMask(int32_t len)
84+ {
85+ uint64_t mask = 0;
86+ uint64_t one = 1;
87+ uint64_t temp = len % FLOAT_VECTOR_SIZE;
88+ for (int64_t i = 0; i < temp; i++) {
89+ mask |= one << i;
90+ }
91+ 
92+ if (len == VECTOR_SIZE) {
93+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
94+ } else if (len >= FLOAT_VECTOR_SIZE) {
95+ AscendC::SetVectorMask<int8_t>(mask, (uint64_t)-1);
96+ } else {
97+ AscendC::SetVectorMask<int8_t>(0x0, mask);
98+ }
99+ }
100+ 
101+ CATLASS_DEVICE
102+ void SetkvSplitCoreNum(uint32_t kvSplitCoreNum_)
103+ {
104+ kvSplitCoreNum = kvSplitCoreNum_;
105+ }
106+ 
107+ CATLASS_DEVICE
108+ void SubCoreCompute(AscendC::GlobalTensor<ElementInput> gInput, AscendC::GlobalTensor<ElementUpdate> gUpdate,
109+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementUpdate> gOCoreTmp,
110+ AscendC::GlobalTensor<ElementUpdate> gl, const LayoutInput &layoutInput,
111+ const LayoutOutput &layoutOutput, const LayoutUpdate &layoutUpdate, uint32_t nIdx,
112+ uint32_t needRowLoop, uint32_t rowLoopIdx, uint32_t rescaleOPingPongFlag,
113+ uint32_t *glFlag, uint32_t taskPingPongFlag)
114+ {
115+ uint32_t curRowNum = layoutInput.shape(0);
116+ uint32_t embed = layoutInput.shape(1);
117+ uint32_t embedRound = layoutInput.stride(0);
118+ uint32_t curRowNumRound = RoundUp<HALF_ELENUM_PER_BLK>(curRowNum);
119+ uint64_t dmUbOffsetCurCycle =
120+ (uint64_t)(rescaleOPingPongFlag * HALF_DM_UB_SIZE + rowLoopIdx * ROW_WISE_CYCLE_TILE);
121+ 
122+ uint32_t oUbOffset = oPingPangFlag * ROW_WISE_CYCLE_TILE * embedRound;
123+ 
124+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(oPingPangFlag + 4);
125+ // *** gl_block = expand_to_block(gl), 存放于 tv
126+ AscendC::DataCopy(goUbTensor32[oUbOffset], gInput,
127+ AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
128+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
129+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
130+ 
131+ AscendC::Abs<float,false>(
132+ tvUbTensor,
133+ goUbTensor32[oUbOffset],
134+ (uint64_t)0,
135+ curRowNum * embedRound / FLOAT_VECTOR_SIZE,
136+ {1, 1, 8, 8}
137+ );
138+ AscendC::PipeBarrier<PIPE_V>();
139+ AscendC::LocalTensor<uint8_t> cmpMaskUb = tvUbTensor.template ReinterpretCast<uint8_t>();
140+ AscendC::CompareScalar(
141+ cmpMaskUb, tvUbTensor, (float)1e10, AscendC::CMPMODE::LE, curRowNum * embedRound
142+ );
143+ AscendC::PipeBarrier<PIPE_V>();
144+ AscendC::Select(
145+ goUbTensor32[oUbOffset], cmpMaskUb, goUbTensor32[oUbOffset], (float)0.0, AscendC::SELMODE::VSEL_TENSOR_SCALAR_MODE, curRowNum * embedRound
146+ );
147+ AscendC::PipeBarrier<PIPE_V>();
148+ 
149+ AscendC::Brcb(tvUbTensor.ReinterpretCast<uint32_t>(),
150+ glUbTensor[taskPingPongFlag].ReinterpretCast<uint32_t>()[rowLoopIdx * ROW_WISE_CYCLE_TILE],
151+ curRowNumRound / FLOAT_BLOCK_SIZE,
152+ AscendC::BrcbRepeatParams(1, 8));
153+ AscendC::PipeBarrier<PIPE_V>();
154+ 
155+ // *** go = go / gl_block
156+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
157+ for (uint32_t vdiv_idx = 0; vdiv_idx < embed / FLOAT_VECTOR_SIZE; ++vdiv_idx) {
158+ AscendC::Div<float, false>(goUbTensor32[oUbOffset + vdiv_idx * FLOAT_VECTOR_SIZE],
159+ goUbTensor32[oUbOffset + vdiv_idx * FLOAT_VECTOR_SIZE],
160+ tvUbTensor, (uint64_t)0,
161+ curRowNum,
162+ AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE,
163+ embedRound / FLOAT_BLOCK_SIZE, 1));
164+ 
165+ }
166+ if (embed % FLOAT_VECTOR_SIZE > 0) {
167+ SetMask(embed % FLOAT_VECTOR_SIZE);
168+ AscendC::Div<float, false>(goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
169+ goUbTensor32[oUbOffset + embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
170+ tvUbTensor,
171+ (uint64_t)0, curRowNum,
172+ AscendC::BinaryRepeatParams(1, 1, 0, embedRound / FLOAT_BLOCK_SIZE,
173+ embedRound / FLOAT_BLOCK_SIZE, 1));
174+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
175+ 
176+ }
177+ AscendC::PipeBarrier<PIPE_V>();
178+ 
179+ // *** go = castfp32to16(go)
180+ if (std::is_same<ElementOutput, bfloat16_t>::value) {
181+ AscendC::Cast<ElementOutput, float, false>(
182+ goUbTensor16[oUbOffset * 2], goUbTensor32[oUbOffset],
183+ AscendC::RoundMode::CAST_RINT, (uint64_t)0,
184+ (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
185+ AscendC::UnaryRepeatParams(1, 1, 4, 8));
186+ } else {
187+ AscendC::Cast<ElementOutput, float, false>(
188+ goUbTensor16[oUbOffset * 2], goUbTensor32[oUbOffset],
189+ AscendC::RoundMode::CAST_NONE, (uint64_t)0,
190+ (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
191+ AscendC::UnaryRepeatParams(1, 1, 4, 8));
192+ }
193+ 
194+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
195+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
196+ 
197+ // ********************* move O to GM ************************
198+ AscendC::DataCopyPad(gOutput, goUbTensor16[oUbOffset * 2],
199+ AscendC::DataCopyExtParams(curRowNum, embed * 2, 0, 0, 0));
200+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(oPingPangFlag + 4);
201+ 
202+ oPingPangFlag = 1 - oPingPangFlag;
203+ }
204+ 
205+ CATLASS_DEVICE
206+ void operator()(AscendC::GlobalTensor<ElementInput> gInput, AscendC::GlobalTensor<ElementUpdate> gUpdate,
207+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementUpdate> gOCoreTmp,
208+ AscendC::GlobalTensor<ElementUpdate> gl, const LayoutInput &layoutInput,
209+ const LayoutUpdate &layoutUpdate, const LayoutOutput &layoutOutput, GemmCoord actualBlockShape,
210+ uint32_t nIdx, uint32_t rescaleOPingPongFlag, uint32_t *glFlag, uint32_t taskPingPongFlag)
211+ {
212+ uint32_t embed = layoutInput.shape(1);
213+ uint32_t rowActual = actualBlockShape.m();
214+ uint32_t columnActual = actualBlockShape.n();
215+ 
216+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
217+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
218+ 
219+ uint32_t curRowSplitSubBlock = rowActual / subBlockNum;
220+ uint32_t rowActualThisSubBlock = (subBlockIdx == 0) ? curRowSplitSubBlock : (rowActual - curRowSplitSubBlock);
221+ uint32_t rowOffsetSubBlock = subBlockIdx * curRowSplitSubBlock;
222+ 
223+ if (rowActualThisSubBlock > 0) {
224+ uint32_t rowLoop = (rowActualThisSubBlock + ROW_WISE_CYCLE_TILE - 1) / ROW_WISE_CYCLE_TILE;
225+ uint32_t needRowLoop = (rowLoop > 1) ? 1 : 0;
226+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoop; rowLoopIdx++) {
227+ uint32_t rowOffsetLoop = rowLoopIdx * ROW_WISE_CYCLE_TILE;
228+ uint32_t rowOffsetCurCycle = rowOffsetSubBlock + rowOffsetLoop;
229+ uint32_t rowActualCurCycle =
230+ (rowLoopIdx == (rowLoop - 1)) ? rowActualThisSubBlock - rowOffsetLoop : ROW_WISE_CYCLE_TILE;
231+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetCurCycle, 0));
232+ auto gInputThisCurCycle = gInput[offsetInput];
233+ auto layoutInputCurCycle = layoutInput.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual));
234+ 
235+ int64_t offsetUpdate = layoutUpdate.GetOffset(MatrixCoord(rowOffsetCurCycle, 0));
236+ auto gUpdateCurCycle = gUpdate[offsetUpdate];
237+ auto layoutUpdateCurCycle = layoutUpdate.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual));
238+ 
239+ int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetCurCycle, 0));
240+ auto gOutputCurCycle = gOutput[offsetOutput];
241+ auto layoutOutputCurCycle = layoutOutput.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual));
242+ 
243+ SubCoreCompute(gInputThisCurCycle, gUpdateCurCycle, gOutputCurCycle,
244+ gOCoreTmp[rowOffsetLoop * embed * kvSplitCoreNum],
245+ gl[rowOffsetLoop * kvSplitCoreNum], layoutInputCurCycle,
246+ layoutOutputCurCycle, layoutUpdateCurCycle, nIdx, needRowLoop, rowLoopIdx,
247+ rescaleOPingPongFlag, glFlag, taskPingPongFlag);
248+ }
249+ }
250+ }
251+ 
252+private:
253+ uint32_t kvSplitCoreNum = 1;
254+ uint32_t oPingPangFlag = 0;
255+ AscendC::LocalTensor<float> loUbTensor;
256+ AscendC::LocalTensor<float> dmUbTensor;
257+ AscendC::LocalTensor<float> glUbTensor[2];
258+ AscendC::LocalTensor<float> tvUbTensor;
259+ AscendC::LocalTensor<ElementOutput> goUbTensor16;
260+ AscendC::LocalTensor<float> goUbTensor32;
261+ AscendC::LocalTensor<float> hmUbTensor;
262+ AscendC::LocalTensor<float> gmUbTensor[2];
263+};
264+ 
265+} // namespace Catlass::Epilogue::Block
266+ 
267+#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_AMLA_TP1_RESCALE_O_HPP
Ainclude/catlass/epilogue/block/block_epilogue_amla_tp1_softmax.hpp+838-0
@@ -0,0 +1,838 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_AMLA_TP1_SOFTMAX_HPP
12+#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_AMLA_TP1_SOFTMAX_HPP
13+ 
14+#include "catlass/catlass.hpp"
15+#include "catlass/arch/cross_core_sync.hpp"
16+#include "catlass/arch/resource.hpp"
17+#include "catlass/epilogue/dispatch_policy.hpp"
18+#include "catlass/epilogue/tile/tile_copy.hpp"
19+#include "catlass/gemm_coord.hpp"
20+#include "catlass/matrix_coord.hpp"
21+ 
22+namespace Catlass::Epilogue::Block {
23+ 
24+template <
25+ class OutputType_,
26+ class InputType_,
27+ class UpdateType_,
28+ class MaskType_>
29+class BlockEpilogue<
30+ EpilogueAtlasA2AMLATP1Softmax,
31+ OutputType_,
32+ InputType_,
33+ UpdateType_,
34+ MaskType_>
35+{
36+public:
37+ // Type aliases
38+ using DispatchPolicy = EpilogueAtlasA2AMLATP1Softmax;
39+ using ArchTag = typename DispatchPolicy::ArchTag;
40+ using ElementOutput = typename OutputType_::Element;
41+ using ElementInput = typename InputType_::Element;
42+ using ElementMask = typename MaskType_::Element;
43+ 
44+ using LayoutOutput = typename OutputType_::Layout;
45+ using LayoutInput = typename InputType_::Layout;
46+ using LayoutMask = typename MaskType_::Layout;
47+ using LayoutUpdate = typename UpdateType_::Layout;
48+ 
49+ using CopyGmToUbInput = Tile::CopyGm2Ub<ArchTag, InputType_>;
50+ using CopyGmToUbMask = Tile::CopyGm2Ub<ArchTag, MaskType_>;
51+ using CopyUbToGmOutput = Tile::CopyUb2Gm<ArchTag, OutputType_>;
52+ 
53+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
54+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
55+ static constexpr uint32_t HALF_VECTOR_SIZE = 128;
56+ static constexpr uint32_t BLOCK_SIZE = 16;
57+ static constexpr uint32_t UB_UINT8_LINE_SIZE = 512;
58+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE_MLA = 16384;
59+ static constexpr uint32_t VECTOR_SIZE = 128;
60+ 
61+ static constexpr uint32_t REDUCE_UB_SIZE = 1024;
62+ static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32;
63+ static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4;
64+ static constexpr uint32_t S_BLOCK_STACK = 4;
65+ static constexpr int64_t UB_FLOAT_LINE_SIZE = 64;
66+ static constexpr uint32_t M_SLICE = 16;
67+ static constexpr uint32_t QK_READY_ID = 1;
68+ static constexpr uint32_t SOFTMAX_READY_ID = 2;
69+ 
70+ static constexpr float LN2 = 0.6931471805599453094172;
71+ static constexpr float RECIP_OF_LN2 = 1 / LN2;
72+ static constexpr float FLOAT_E_SCALAR = 8388608;
73+ static constexpr uint32_t PRE_AUTOADD_READY_ID = 0;
74+ static constexpr uint32_t SM_AUTOADD_READY_ID = 7;
75+ static constexpr uint32_t BRCB_SIZE = 128;
76+ 
77+ CATLASS_DEVICE
78+ BlockEpilogue(Arch::Resource<ArchTag> &resource, half tor_, uint32_t kvSplitCoreNum_ = 1)
79+ {
80+ // Allocate UB space
81+ constexpr uint32_t LS_UB_TENSOR_OFFSET = 0;
82+ constexpr uint32_t LP_UB_TENSOR_OFFSET = 0;
83+ constexpr uint32_t LM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA;
84+ constexpr uint32_t HM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE;
85+ constexpr uint32_t DM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 2 * UB_UINT8_LINE_SIZE;
86+ constexpr uint32_t LL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 4 * UB_UINT8_LINE_SIZE;
87+ constexpr uint32_t GM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 5 * UB_UINT8_LINE_SIZE;
88+ constexpr uint32_t GL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 7 * UB_UINT8_LINE_SIZE;
89+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA;
90+ 
91+ constexpr uint32_t NBRCB_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE_MLA;
92+ constexpr uint32_t HN_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 9 * UB_UINT8_LINE_SIZE;
93+ constexpr uint32_t DN_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 10 * UB_UINT8_LINE_SIZE;
94+ constexpr uint32_t GN_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 12 * UB_UINT8_LINE_SIZE;
95+ constexpr uint32_t S32_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 14 * UB_UINT8_LINE_SIZE;
96+ constexpr uint32_t S16Tmp_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 16 * UB_UINT8_LINE_SIZE;
97+ constexpr uint32_t S16_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 18 * UB_UINT8_LINE_SIZE;
98+ constexpr uint32_t HC_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 20 * UB_UINT8_LINE_SIZE;
99+ constexpr uint32_t GC_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 22 * UB_UINT8_LINE_SIZE;
100+ constexpr uint32_t EPS_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 24 * UB_UINT8_LINE_SIZE;
101+ constexpr uint32_t N_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 26 * UB_UINT8_LINE_SIZE;
102+ 
103+ tor = tor_;
104+ kvSplitCoreNum = kvSplitCoreNum_;
105+ lsUbTensor = resource.ubBuf.template GetBufferByByte<float>(LS_UB_TENSOR_OFFSET);
106+ lpUbTensor = resource.ubBuf.template GetBufferByByte<ElementOutput>(LP_UB_TENSOR_OFFSET);
107+ lmUbTensor = resource.ubBuf.template GetBufferByByte<float>(LM_UB_TENSOR_OFFSET);
108+ hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);
109+ gmUbTensor[0] = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
110+ gmUbTensor[1] = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET + UB_UINT8_LINE_SIZE);
111+ dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
112+ llUbTensor = resource.ubBuf.template GetBufferByByte<float>(LL_UB_TENSOR_OFFSET);
113+ tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
114+ glUbTensor[0] = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
115+ glUbTensor[1] = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET + UB_UINT8_LINE_SIZE);
116+ 
117+ hnUbTensor = resource.ubBuf.template GetBufferByByte<float>(HN_UB_TENSOR_OFFSET);
118+ dnUbTensor = resource.ubBuf.template GetBufferByByte<float>(DN_UB_TENSOR_OFFSET);
119+ gnUbTensor = resource.ubBuf.template GetBufferByByte<float>(GN_UB_TENSOR_OFFSET);
120+ s32UbTensor = resource.ubBuf.template GetBufferByByte<float>(S32_UB_TENSOR_OFFSET);
121+ s16UbTensorTmp = resource.ubBuf.template GetBufferByByte<ElementOutput>(S16Tmp_UB_TENSOR_OFFSET);
122+ NBrcbUbTensor = resource.ubBuf.template GetBufferByByte<int32_t>(NBRCB_UB_TENSOR_OFFSET);
123+ s16UbTensor = resource.ubBuf.template GetBufferByByte<float>(S16_UB_TENSOR_OFFSET);
124+ hcUbTensor = resource.ubBuf.template GetBufferByByte<float>(HC_UB_TENSOR_OFFSET);
125+ gcUbTensor = resource.ubBuf.template GetBufferByByte<float>(GC_UB_TENSOR_OFFSET);
126+ epsUbTensor = resource.ubBuf.template GetBufferByByte<float>(EPS_UB_TENSOR_OFFSET);
127+ NUbTensor = resource.ubBuf.template GetBufferByByte<int32_t>(N_UB_TENSOR_OFFSET);
128+ }
129+ 
130+ CATLASS_DEVICE
131+ ~BlockEpilogue() {}
132+ 
133+ CATLASS_DEVICE
134+ void SetVecMask(int32_t len)
135+ {
136+ uint64_t mask = 0;
137+ uint64_t one = 1;
138+ uint64_t temp = len % FLOAT_VECTOR_SIZE;
139+ for (int64_t i = 0; i < temp; i++) {
140+ mask |= one << i;
141+ }
142+ 
143+ if (len == VECTOR_SIZE || len == 0) {
144+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
145+ } else if (len >= FLOAT_VECTOR_SIZE) {
146+ AscendC::SetVectorMask<int8_t>(mask, (uint64_t)-1);
147+ } else {
148+ AscendC::SetVectorMask<int8_t>(0x0, mask);
149+ }
150+ }
151+ 
152+ CATLASS_DEVICE
153+ void SetkvSplitCoreNum(uint32_t kvSplitCoreNum_)
154+ {
155+ kvSplitCoreNum = kvSplitCoreNum_;
156+ }
157+ 
158+ CATLASS_DEVICE
159+ void SetBlockReduceMask(int32_t len)
160+ {
161+ if (len > 8 || len < 1) {
162+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
163+ return;
164+ }
165+ uint64_t subMask = ((uint64_t)1 << len) - 1;
166+ uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask + (subMask << 56) +
167+ (subMask << 40) + (subMask << 24) + (subMask << 8);
168+ AscendC::SetVectorMask<int8_t>(maskValue, maskValue);
169+ }
170+ 
171+ CATLASS_DEVICE
172+ void RowsumSPECTILE512(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowsumUb,
173+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
174+ uint32_t numElemsAligned)
175+ {
176+ AscendC::BlockReduceSum<float, false>(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0,
177+ 1, 1, 8);
178+ AscendC::PipeBarrier<PIPE_V>();
179+ 
180+ AscendC::BlockReduceSum<float, false>(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
181+ numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0,
182+ 1, 1, 8);
183+ AscendC::PipeBarrier<PIPE_V>();
184+ AscendC::BlockReduceSum<float, false>(rowsumUb, tvUbTensor[REDUCE_UB_SIZE],
185+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0,
186+ 1, 1, 8);
187+ AscendC::PipeBarrier<PIPE_V>();
188+ }
189+ 
190+ CATLASS_DEVICE
191+ void RowsumSPECTILE256(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowsumUb,
192+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
193+ uint32_t numElemsAligned)
194+ {
195+ AscendC::BlockReduceSum<float, false>(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0,
196+ 1, 1, 8);
197+ AscendC::PipeBarrier<PIPE_V>();
198+ SetVecMask(ROW_OPS_SPEC_MASK_32);
199+ AscendC::BlockReduceSum<float, false>(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4);
200+ AscendC::PipeBarrier<PIPE_V>();
201+ SetBlockReduceMask(ROW_OPS_SPEC_MASK_4);
202+ AscendC::BlockReduceSum<float, false>(
203+ rowsumUb, tvUbTensor[REDUCE_UB_SIZE],
204+ (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
205+ AscendC::PipeBarrier<PIPE_V>();
206+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
207+ }
208+ 
209+ CATLASS_DEVICE
210+ void RowsumTAILTILE(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowsumUb,
211+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
212+ uint32_t numElemsAligned)
213+ {
214+ if (numElems >= FLOAT_VECTOR_SIZE) {
215+ AscendC::BlockReduceSum<float, false>(tvUbTensor, srcUb, numRowsRound, 0, 1, 1,
216+ numElemsAligned / FLOAT_BLOCK_SIZE);
217+ AscendC::PipeBarrier<PIPE_V>();
218+ AscendC::BlockReduceSum<float, false>(
219+ rowsumUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0,
220+ 1, 1, 8);
221+ AscendC::PipeBarrier<PIPE_V>();
222+ for (uint64_t rowSumIdx = 1; rowSumIdx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowSumIdx) {
223+ AscendC::BlockReduceSum<float, false>(tvUbTensor, srcUb[rowSumIdx * FLOAT_VECTOR_SIZE], numRowsRound, 0,
224+ 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
225+ AscendC::PipeBarrier<PIPE_V>();
226+ AscendC::BlockReduceSum<float, false>(
227+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
228+ (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
229+ AscendC::PipeBarrier<PIPE_V>();
230+ SetVecMask(numRowsRound);
231+ AscendC::Add<float, false>(rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
232+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
233+ AscendC::PipeBarrier<PIPE_V>();
234+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
235+ }
236+ }
237+ if (numElems % FLOAT_VECTOR_SIZE > 0) {
238+ SetVecMask(numElems % FLOAT_VECTOR_SIZE);
239+ AscendC::BlockReduceSum<float, false>(tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
240+ numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
241+ AscendC::PipeBarrier<PIPE_V>();
242+ SetBlockReduceMask((numElems % FLOAT_VECTOR_SIZE + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE);
243+ if (numElems < FLOAT_VECTOR_SIZE) {
244+ AscendC::BlockReduceSum<float, false>(
245+ rowsumUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
246+ 0, 1, 1, 8);
247+ AscendC::PipeBarrier<PIPE_V>();
248+ } else {
249+ AscendC::BlockReduceSum<float, false>(
250+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
251+ (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
252+ AscendC::PipeBarrier<PIPE_V>();
253+ SetVecMask(numRowsRound);
254+ AscendC::Add<float, false>(rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
255+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
256+ AscendC::PipeBarrier<PIPE_V>();
257+ }
258+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
259+ }
260+ }
261+ 
262+ CATLASS_DEVICE
263+ void RowmaxSPECTILE512(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowmaxUb,
264+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
265+ uint32_t numElemsAligned)
266+ {
267+ AscendC::BlockReduceMax<float, false>(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0,
268+ 1, 1, 8);
269+ AscendC::PipeBarrier<PIPE_V>();
270+ AscendC::BlockReduceMax<float, false>(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
271+ numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0,
272+ 1, 1, 8);
273+ AscendC::PipeBarrier<PIPE_V>();
274+ AscendC::BlockReduceMax<float, false>(rowmaxUb, tvUbTensor[REDUCE_UB_SIZE],
275+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0,
276+ 1, 1, 8);
277+ AscendC::PipeBarrier<PIPE_V>();
278+ }
279+ 
280+ CATLASS_DEVICE
281+ void RowmaxSPECTILE256(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowmaxUb,
282+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
283+ uint32_t numElemsAligned)
284+ {
285+ AscendC::BlockReduceMax<float, false>(tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0,
286+ 1, 1, 8);
287+ AscendC::PipeBarrier<PIPE_V>();
288+ SetVecMask(ROW_OPS_SPEC_MASK_32);
289+ AscendC::BlockReduceMax<float, false>(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4);
290+ AscendC::PipeBarrier<PIPE_V>();
291+ SetBlockReduceMask(ROW_OPS_SPEC_MASK_4);
292+ AscendC::BlockReduceMax<float, false>(
293+ rowmaxUb, tvUbTensor[REDUCE_UB_SIZE],
294+ (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
295+ AscendC::PipeBarrier<PIPE_V>();
296+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
297+ }
298+ 
299+ CATLASS_DEVICE
300+ void RowmaxTAILTILE(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowmaxUb,
301+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
302+ uint32_t numElemsAligned)
303+ {
304+ if (numElems >= FLOAT_VECTOR_SIZE) {
305+ AscendC::BlockReduceMax<float, false>(tvUbTensor, srcUb, numRowsRound, 0, 1, 1,
306+ numElemsAligned / FLOAT_BLOCK_SIZE);
307+ AscendC::PipeBarrier<PIPE_V>();
308+ AscendC::BlockReduceMax<float, false>(
309+ rowmaxUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0,
310+ 1, 1, 8);
311+ AscendC::PipeBarrier<PIPE_V>();
312+ for (uint64_t rowmax_idx = 1; rowmax_idx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowmax_idx) {
313+ AscendC::BlockReduceMax<float, false>(tvUbTensor, srcUb[rowmax_idx * FLOAT_VECTOR_SIZE], numRowsRound,
314+ 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
315+ AscendC::PipeBarrier<PIPE_V>();
316+ AscendC::BlockReduceMax<float, false>(
317+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
318+ (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
319+ AscendC::PipeBarrier<PIPE_V>();
320+ SetVecMask(numRowsRound);
321+ AscendC::Max<float, false>(rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
322+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
323+ AscendC::PipeBarrier<PIPE_V>();
324+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
325+ }
326+ }
327+ if (numElems % FLOAT_VECTOR_SIZE > 0) {
328+ SetVecMask(numElems % FLOAT_VECTOR_SIZE);
329+ AscendC::BlockReduceMax<float, false>(tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
330+ numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
331+ AscendC::PipeBarrier<PIPE_V>();
332+ SetBlockReduceMask((numElems % FLOAT_VECTOR_SIZE + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE);
333+ if (numElems < FLOAT_VECTOR_SIZE) {
334+ AscendC::BlockReduceMax<float, false>(
335+ rowmaxUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
336+ 0, 1, 1, 8);
337+ AscendC::PipeBarrier<PIPE_V>();
338+ } else {
339+ AscendC::BlockReduceMax<float, false>(
340+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
341+ (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
342+ AscendC::PipeBarrier<PIPE_V>();
343+ SetVecMask(numRowsRound);
344+ AscendC::Max<float, false>(rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
345+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
346+ AscendC::PipeBarrier<PIPE_V>();
347+ }
348+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
349+ }
350+ }
351+ 
352+ CATLASS_DEVICE
353+ void SubCoreCompute(AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
354+ uint32_t m, uint32_t nReal, uint32_t nStride, uint32_t pingpongFlag, uint32_t rowOffset,
355+ uint32_t sUbOffset, uint32_t nIdx, uint32_t isLastNTile, uint32_t *glFlag,
356+ uint32_t taskPingPongFlag, uint32_t gSPingPongFlag, uint32_t isLastM, uint32_t isFirstM)
357+ {
358+ uint32_t round_m = (m + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE * FLOAT_BLOCK_SIZE;
359+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(pingpongFlag);
360+ // input QK
361+ AscendC::DataCopy(lsUbTensor[sUbOffset], gInput, AscendC::DataCopyParams(m, nStride / FLOAT_BLOCK_SIZE, 0, 0));
362+ 
363+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
364+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
365+ 
366+ // *** ls = tor * ls
367+ AscendC::Muls<float, false>(lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], tor, (uint64_t)0,
368+ (m * nStride + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
369+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
370+ 
371+ AscendC::PipeBarrier<PIPE_V>();
372+ 
373+ 
374+ if (nReal == 512) {
375+ RowmaxSPECTILE512(lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride);
376+ } else if (nReal == 256) {
377+ RowmaxSPECTILE256(lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride);
378+ } else {
379+ RowmaxTAILTILE(lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride);
380+ }
381+ 
382+ 
383+ if (nIdx == 0) {
384+ AscendC::DataCopy(hmUbTensor[rowOffset], lmUbTensor[rowOffset],
385+ AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0));
386+ 
387+ AscendC::PipeBarrier<PIPE_V>();
388+ } else {
389+ SetVecMask(m);
390+ // *** hm = vmax(lm, gm)
391+ AscendC::Max<float, false>(hmUbTensor[rowOffset], lmUbTensor[rowOffset], gmUbTensor[taskPingPongFlag][rowOffset], (uint64_t)0,
392+ 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
393+ 
394+ AscendC::PipeBarrier<PIPE_V>();
395+ // *** dm = gm - hm
396+ AscendC::Sub<float, false>(dmUbTensor[rowOffset],
397+ gmUbTensor[taskPingPongFlag][rowOffset], hmUbTensor[rowOffset], (uint64_t)0, 1,
398+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
399+ 
400+ AscendC::PipeBarrier<PIPE_V>();
401+ // *** dm = exp(dm)
402+ AscendC::Exp<float, false>(dmUbTensor[rowOffset],
403+ dmUbTensor[rowOffset],
404+ (uint64_t)0, 1, AscendC::UnaryRepeatParams(1, 1, 8, 8));
405+ }
406+ 
407+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
408+ AscendC::PipeBarrier<PIPE_V>();
409+ // *** gm = hm
410+ AscendC::DataCopy(gmUbTensor[taskPingPongFlag][rowOffset], hmUbTensor[rowOffset],
411+ AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0));
412+ AscendC::PipeBarrier<PIPE_V>();
413+ 
414+ // ni = hnUbTensor = mi / -LN2
415+ SetVecMask(m);
416+ AscendC::Muls<float, false>(
417+ hnUbTensor[rowOffset],
418+ hmUbTensor[rowOffset],
419+ (-1.0f) * RECIP_OF_LN2,
420+ (uint64_t)0,
421+ 1,
422+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
423+ AscendC::PipeBarrier<PIPE_V>();
424+ AscendC::Cast<float, float, false>(
425+ hnUbTensor[rowOffset], hnUbTensor[rowOffset], AscendC::RoundMode::CAST_ROUND, (uint64_t)0,
426+ 1, AscendC::UnaryRepeatParams(1, 1, 8, 8));
427+ AscendC::PipeBarrier<PIPE_V>();
428+ 
429+ // dnUbTensor = ni - ni-1
430+ if (nIdx != 0) {
431+ AscendC::Sub<float, false>(
432+ dnUbTensor[rowOffset],
433+ hnUbTensor[rowOffset],
434+ gnUbTensor[rowOffset],
435+ (uint64_t)0,
436+ 1,
437+ AscendC::BinaryRepeatParams(
438+ 1, 1, 1, 8, 8, 8));
439+ AscendC::PipeBarrier<PIPE_V>();
440+ } else {
441+ // n1 = n0, dn = 0
442+ AscendC::Duplicate<float, false>(dnUbTensor[rowOffset], 0.0f, (uint64_t)0, 1, 1, 8);
443+ AscendC::PipeBarrier<PIPE_V>();
444+ }
445+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
446+ AscendC::DataCopy(gnUbTensor[rowOffset], hnUbTensor[rowOffset], AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0));
447+ AscendC::PipeBarrier<PIPE_V>();
448+ 
449+ // mi / LN2
450+ SetVecMask(m);
451+ AscendC::Muls<float, false>(
452+ tvUbTensor,
453+ hmUbTensor[rowOffset],
454+ RECIP_OF_LN2,
455+ (uint64_t)0,
456+ 1,
457+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
458+ AscendC::PipeBarrier<PIPE_V>();
459+ // hnUbTensor = ni + mi / LN2
460+ AscendC::Add<float, false>(
461+ hnUbTensor[rowOffset],
462+ hnUbTensor[rowOffset],
463+ tvUbTensor,
464+ (uint64_t)0,
465+ 1,
466+ AscendC::BinaryRepeatParams(
467+ 1, 1, 1, 8, 8, 8));
468+ AscendC::PipeBarrier<PIPE_V>();
469+ // ln2 * (ni + mi / ln2)
470+ AscendC::Muls<float, false>(
471+ hnUbTensor[rowOffset],
472+ hnUbTensor[rowOffset],
473+ LN2,
474+ (uint64_t)0,
475+ 1,
476+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
477+ AscendC::PipeBarrier<PIPE_V>();
478+ // s32 = exp(ln2 * (ni + mi / ln2))
479+ AscendC::Exp<float, false>(
480+ s32UbTensor[rowOffset],
481+ hnUbTensor[rowOffset],
482+ (uint64_t)0,
483+ 1,
484+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
485+ AscendC::PipeBarrier<PIPE_V>();
486+ // s16 = s32 -> s16
487+ if (std::is_same<ElementOutput, bfloat16_t>::value) {
488+ AscendC::Cast<ElementOutput, float, false>(
489+ s16UbTensorTmp[rowOffset], s32UbTensor[rowOffset], AscendC::RoundMode::CAST_ROUND, (uint64_t)0,
490+ 1, AscendC::UnaryRepeatParams(1, 1, 4, 8));
491+ } else {
492+ AscendC::Cast<ElementOutput, float, false>(
493+ s16UbTensorTmp[rowOffset], s32UbTensor[rowOffset], AscendC::RoundMode::CAST_RINT, (uint64_t)0,
494+ 1, AscendC::UnaryRepeatParams(1, 1, 4, 8));
495+ }
496+ AscendC::PipeBarrier<PIPE_V>();
497+ // s16 = s16 -> s16
498+ AscendC::Cast<float, ElementOutput, false>(
499+ s16UbTensor[rowOffset], s16UbTensorTmp[rowOffset], AscendC::RoundMode::CAST_NONE, (uint64_t)0,
500+ 1, AscendC::UnaryRepeatParams(1, 1, 8, 4));
501+ AscendC::PipeBarrier<PIPE_V>();
502+ 
503+ // *** hm_block = expand_to_block(hm), 存放于 tv
504+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
505+ AscendC::Brcb(tvUbTensor.template ReinterpretCast<uint32_t>(),
506+ hmUbTensor[rowOffset].template ReinterpretCast<uint32_t>(), round_m / FLOAT_BLOCK_SIZE,
507+ AscendC::BrcbRepeatParams(1, 8));
508+ AscendC::PipeBarrier<PIPE_V>();
509+ 
510+ // *** ls = ls - hm_block
511+ for (uint32_t subIdx = 0; subIdx < nReal / FLOAT_VECTOR_SIZE; ++subIdx) {
512+ AscendC::Sub<float, false>(
513+ lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE],
514+ tvUbTensor, (uint64_t)0, m,
515+ AscendC::BinaryRepeatParams(1, 1, 0, nStride / FLOAT_BLOCK_SIZE, nStride / FLOAT_BLOCK_SIZE, 1));
516+ }
517+ if (nReal % FLOAT_VECTOR_SIZE > 0) {
518+ SetVecMask(nReal % FLOAT_VECTOR_SIZE);
519+ AscendC::Sub<float, false>(
520+ lsUbTensor[sUbOffset][nReal / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
521+ lsUbTensor[sUbOffset][nReal / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0, m,
522+ AscendC::BinaryRepeatParams(1, 1, 0, nStride / FLOAT_BLOCK_SIZE, nStride / FLOAT_BLOCK_SIZE, 1));
523+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
524+ }
525+ AscendC::PipeBarrier<PIPE_V>();
526+ 
527+ // *** ls = exp(ls)
528+ AscendC::Exp<float, false>(lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], (uint64_t)0,
529+ (m * nStride + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
530+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
531+ AscendC::PipeBarrier<PIPE_V>();
532+ 
533+ // *** ll = rowsum(ls32)
534+ if (nReal == 512) {
535+ RowsumSPECTILE512(lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride);
536+ } else if (nReal == 256) {
537+ RowsumSPECTILE256(lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride);
538+ } else {
539+ RowsumTAILTILE(lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride);
540+ }
541+ 
542+ AscendC::Brcb(
543+ tvUbTensor.ReinterpretCast<uint32_t>(),
544+ s16UbTensor.ReinterpretCast<uint32_t>()[rowOffset],
545+ round_m / FLOAT_BLOCK_SIZE,
546+ AscendC::BrcbRepeatParams(1, 8));
547+ AscendC::PipeBarrier<PIPE_V>();
548+ 
549+ // *** p = p * s16
550+ for (uint32_t subIdx = 0; subIdx < nReal / FLOAT_VECTOR_SIZE; ++subIdx) {
551+ AscendC::Mul<float, false>(
552+ lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE],
553+ tvUbTensor, (uint64_t)0, m,
554+ AscendC::BinaryRepeatParams(1, 1, 0, nStride / FLOAT_BLOCK_SIZE, nStride / FLOAT_BLOCK_SIZE, 1));
555+ 
556+ }
557+ if (nReal % FLOAT_VECTOR_SIZE > 0) {
558+ SetVecMask(nReal % FLOAT_VECTOR_SIZE);
559+ AscendC::Mul<float, false>(
560+ lsUbTensor[sUbOffset][nReal / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
561+ lsUbTensor[sUbOffset][nReal / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0, m,
562+ AscendC::BinaryRepeatParams(1, 1, 0, nStride / FLOAT_BLOCK_SIZE, nStride / FLOAT_BLOCK_SIZE, 1));
563+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
564+ }
565+ AscendC::PipeBarrier<PIPE_V>();
566+ 
567+ // *** lp = castfp32to16(ls)
568+ if (std::is_same<ElementOutput, bfloat16_t>::value) {
569+ AscendC::Cast<ElementOutput, float, false>(
570+ lpUbTensor[sUbOffset * 2], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_RINT, (uint64_t)0,
571+ (m * nStride + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, AscendC::UnaryRepeatParams(1, 1, 4, 8));
572+ } else {
573+ AscendC::Cast<ElementOutput, float, false>(
574+ lpUbTensor[sUbOffset * 2], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_NONE, (uint64_t)0,
575+ (m * nStride + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, AscendC::UnaryRepeatParams(1, 1, 4, 8));
576+ }
577+ 
578+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
579+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
580+ 
581+ AscendC::DataCopy(gOutput, lpUbTensor[sUbOffset * 2], AscendC::DataCopyParams(m, nStride * 2 / 32, 0, 0));
582+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(pingpongFlag);
583+ 
584+ if (isLastM) {
585+ Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxReady);
586+ }
587+ if (nIdx == 0) {
588+ // *** gl = ll
589+ AscendC::DataCopy(glUbTensor[taskPingPongFlag][rowOffset], llUbTensor[rowOffset],
590+ AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0));
591+ 
592+ AscendC::PipeBarrier<PIPE_V>();
593+ } else {
594+ SetVecMask(m);
595+ // // *** gl = dm * gl
596+ AscendC::Mul<float, false>(
597+ glUbTensor[taskPingPongFlag][rowOffset], dmUbTensor[rowOffset],
598+ glUbTensor[taskPingPongFlag][rowOffset], (uint64_t)0, 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
599+ AscendC::PipeBarrier<PIPE_V>();
600+ 
601+ // // *** gl = ll + gl
602+ AscendC::Add<float, false>(glUbTensor[taskPingPongFlag][rowOffset], glUbTensor[taskPingPongFlag][rowOffset], llUbTensor[rowOffset], (uint64_t)0,
603+ 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
604+ AscendC::PipeBarrier<PIPE_V>();
605+ 
606+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
607+ }
608+ // -----------------------------------------
609+ SetVecMask(m);
610+ if (isLastNTile) {
611+ AscendC::Mul<float, false>(
612+ glUbTensor[taskPingPongFlag][rowOffset],
613+ glUbTensor[taskPingPongFlag][rowOffset],
614+ s16UbTensor[rowOffset],
615+ (uint64_t)0,
616+ 1,
617+ AscendC::BinaryRepeatParams(
618+ 1, 1, 1, 8, 8, 8));
619+ AscendC::PipeBarrier<PIPE_V>();
620+ }
621+ 
622+ if (nIdx > 0) {
623+ AscendC::Div<float, false>(
624+ hcUbTensor[rowOffset],
625+ s32UbTensor[rowOffset],
626+ s16UbTensor[rowOffset],
627+ (uint64_t)0,
628+ 1,
629+ AscendC::BinaryRepeatParams(
630+ 1, 1, 1, 8, 8, 8));
631+ AscendC::PipeBarrier<PIPE_V>();
632+ }
633+ // ci = s32 / s16
634+ if (nIdx == 0) {
635+ AscendC::Duplicate<float, false>(gcUbTensor[rowOffset], 1.0, (uint64_t)0, 1, 1, 8);
636+ AscendC::PipeBarrier<PIPE_V>();
637+ }
638+ 
639+ if (nIdx > 0) {
640+ // eps = ci-1 / ci
641+ AscendC::Div<float, false>(
642+ epsUbTensor[rowOffset],
643+ gcUbTensor[rowOffset],
644+ hcUbTensor[rowOffset],
645+ (uint64_t)0,
646+ 1,
647+ AscendC::BinaryRepeatParams(
648+ 1, 1, 1, 8, 8, 8));
649+ AscendC::PipeBarrier<PIPE_V>();
650+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
651+ AscendC::DataCopy(gcUbTensor[rowOffset], hcUbTensor[rowOffset], AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0));
652+ AscendC::PipeBarrier<PIPE_V>();
653+ // eps = eps - 1.0
654+ SetVecMask(m);
655+ AscendC::Adds<float, false>(
656+ epsUbTensor[rowOffset],
657+ epsUbTensor[rowOffset],
658+ (float)(-1.0f),
659+ (uint64_t)0,
660+ 1,
661+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
662+ AscendC::PipeBarrier<PIPE_V>();
663+ // eps = eps * 1.5
664+ AscendC::Muls<float, false>(
665+ epsUbTensor[rowOffset],
666+ epsUbTensor[rowOffset],
667+ (float)(1.5f),
668+ (uint64_t)0,
669+ 1,
670+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
671+ AscendC::PipeBarrier<PIPE_V>();
672+ // max((ni - ni-1), -30.0)
673+ AscendC::Maxs<float, false>(
674+ dnUbTensor[rowOffset],
675+ dnUbTensor[rowOffset],
676+ (float)(-30.0f),
677+ (uint64_t)0,
678+ 1,
679+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
680+ AscendC::PipeBarrier<PIPE_V>();
681+ // eps = eps + 10^-6
682+ AscendC::Adds<float, false>(
683+ epsUbTensor[rowOffset],
684+ epsUbTensor[rowOffset],
685+ (float)(0.000001f),
686+ (uint64_t)0,
687+ 1,
688+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
689+ AscendC::PipeBarrier<PIPE_V>();
690+ // eps + 10^-6 + max((ni - ni-1), -30.0)
691+ AscendC::Add<float, false>(
692+ epsUbTensor[rowOffset],
693+ epsUbTensor[rowOffset],
694+ dnUbTensor[rowOffset],
695+ (uint64_t)0,
696+ 1,
697+ AscendC::BinaryRepeatParams(
698+ 1, 1, 1, 8, 8, 8));
699+ AscendC::PipeBarrier<PIPE_V>();
700+ // eps = eps * 2e23
701+ AscendC::Muls<float, false>(
702+ epsUbTensor[rowOffset],
703+ epsUbTensor[rowOffset],
704+ FLOAT_E_SCALAR,
705+ (uint64_t)0,
706+ 1,
707+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
708+ AscendC::PipeBarrier<PIPE_V>();
709+ // N
710+ AscendC::Cast<int32_t, float, false>(
711+ NUbTensor[gSPingPongFlag * UB_FLOAT_LINE_SIZE + rowOffset], epsUbTensor[rowOffset], AscendC::RoundMode::CAST_ROUND, (uint64_t)0,
712+ 1, AscendC::UnaryRepeatParams(1, 1, 8, 8));
713+ AscendC::PipeBarrier<PIPE_V>();
714+ }
715+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
716+ }
717+ CATLASS_DEVICE
718+ void AmlaAutoadd(AscendC::GlobalTensor<int32_t> gOTmpInt32,
719+ uint32_t subM, uint32_t embed, uint32_t gSPingPongFlag
720+ )
721+ {
722+ // update OTmp
723+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID5);
724+ for (uint32_t i = 0; i < BRCB_SIZE / FLOAT_BLOCK_SIZE; i++){
725+ AscendC::Brcb(
726+ NBrcbUbTensor.ReinterpretCast<uint32_t>()[i * FLOAT_BLOCK_SIZE],
727+ NUbTensor.ReinterpretCast<uint32_t>()[gSPingPongFlag * UB_FLOAT_LINE_SIZE],
728+ subM / FLOAT_BLOCK_SIZE,
729+ AscendC::BrcbRepeatParams(BRCB_SIZE / FLOAT_BLOCK_SIZE, BRCB_SIZE));
730+ }
731+ AscendC::PipeBarrier<PIPE_V>();
732+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID5);
733+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID5);
734+ 
735+ Arch::CrossCoreWaitFlag(preAutoAddReady);
736+ 
737+ AscendC::SetAtomicAdd<int32_t>();
738+ 
739+ for (uint32_t i = 0; i < embed / BRCB_SIZE; i++) {
740+ AscendC::DataCopy(
741+ gOTmpInt32[i * BRCB_SIZE],
742+ NBrcbUbTensor,
743+ AscendC::DataCopyParams(
744+ subM, BRCB_SIZE / FLOAT_BLOCK_SIZE, 0, (embed - BRCB_SIZE) / FLOAT_BLOCK_SIZE));
745+ }
746+ 
747+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID5);
748+ AscendC::SetAtomicNone();
749+ }
750+ 
751+ CATLASS_DEVICE
752+ void operator()(AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
753+ AscendC::GlobalTensor<int32_t> gOTmpInt32,
754+ const LayoutOutput &layoutOutput, const LayoutInput &layoutInput,
755+ const LayoutUpdate &layoutOTmp,
756+ GemmCoord actualBlockShape,
757+ uint32_t nIdx, uint32_t isLastNTile, uint32_t *glFlag, uint32_t taskPingPongFlag, uint32_t gSPingPongFlag)
758+ {
759+ uint32_t cur_head_num = actualBlockShape.m();
760+ uint32_t qkN = actualBlockShape.n();
761+ uint32_t embed = actualBlockShape.k();
762+ uint32_t qkRoundN = layoutInput.stride(0);
763+ 
764+ uint32_t pingpongFlag = 0;
765+ 
766+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
767+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
768+ 
769+ uint32_t subM = (subBlockIdx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2;
770+ 
771+ uint32_t mEnd = (subM + M_SLICE - 1) / M_SLICE;
772+ 
773+ for (uint32_t mInd = 0; mInd < mEnd; mInd++) {
774+ uint32_t rowOffset = mInd * M_SLICE;
775+ uint32_t currM = mInd == mEnd - 1 ? subM - rowOffset : M_SLICE;
776+ uint32_t isLastM = mInd == mEnd - 1 ? 1 : 0;
777+ uint32_t isFirstM = mInd == 0 ? 1 : 0;
778+ uint32_t sUbOffset = pingpongFlag * 8192;
779+ int64_t offsetOutput = rowOffset * qkRoundN;
780+ auto gOutputThisSubBlock = gOutput[offsetOutput];
781+ int64_t offsetInput = rowOffset * qkRoundN;
782+ auto gInputThisSubBlock = gInput[offsetInput];
783+ if (mInd == 0) {
784+ Arch::CrossCoreWaitFlag(qkReady);
785+ }
786+ if (currM == 0) {
787+ continue;
788+ }
789+ SubCoreCompute(gOutputThisSubBlock, gInputThisSubBlock, currM, qkN, qkRoundN, pingpongFlag, rowOffset,
790+ sUbOffset, nIdx, isLastNTile, glFlag, taskPingPongFlag, gSPingPongFlag, isLastM, isFirstM);
791+ pingpongFlag = 1 - pingpongFlag;
792+ }
793+ if (nIdx > 0) {
794+ int64_t gmOffsetOTmp = layoutOTmp.GetOffset(MatrixCoord(subM, 0));
795+ auto gOTmpInt32ThisSubBlock = gOTmpInt32[gmOffsetOTmp];
796+ AmlaAutoadd(gOTmpInt32, subM, embed, gSPingPongFlag);
797+ }
798+ Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxAutoAddReady);
799+ }
800+ 
801+private:
802+ float tor;
803+ uint32_t pingpongFlag = 0;
804+ uint32_t kvSplitCoreNum = 1;
805+ 
806+ AscendC::LocalTensor<float> lsUbTensor;
807+ AscendC::LocalTensor<ElementOutput> lpUbTensor;
808+ AscendC::LocalTensor<float> lmUbTensor;
809+ AscendC::LocalTensor<float> hmUbTensor;
810+ AscendC::LocalTensor<float> gmUbTensor[2];
811+ AscendC::LocalTensor<float> dmUbTensor;
812+ AscendC::LocalTensor<float> llUbTensor;
813+ AscendC::LocalTensor<float> tvUbTensor;
814+ AscendC::LocalTensor<float> glUbTensor[2];
815+ 
816+ // new add
817+ AscendC::LocalTensor<float> hnUbTensor;
818+ AscendC::LocalTensor<float> dnUbTensor;
819+ AscendC::LocalTensor<float> gnUbTensor;
820+ AscendC::LocalTensor<float> s32UbTensor;
821+ AscendC::LocalTensor<ElementOutput> s16UbTensorTmp;
822+ AscendC::LocalTensor<float> s16UbTensor;
823+ AscendC::LocalTensor<float> hcUbTensor;
824+ AscendC::LocalTensor<float> gcUbTensor;
825+ AscendC::LocalTensor<float> epsUbTensor;
826+ AscendC::LocalTensor<int32_t> NUbTensor;
827+ AscendC::LocalTensor<int32_t> NBrcbUbTensor;
828+ Arch::CrossCoreFlag preAutoAddReady{PRE_AUTOADD_READY_ID};
829+ Arch::CrossCoreFlag softmaxAutoAddReady{SM_AUTOADD_READY_ID};
830+ 
831+ Arch::CrossCoreFlag qkReady{QK_READY_ID};
832+ Arch::CrossCoreFlag softmaxReady{SOFTMAX_READY_ID};
833+ 
834+};
835+ 
836+} // namespace Catlass::Epilogue::Block
837+ 
838+#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_AMLA_TP1_SOFTMAX_HPP
Minclude/catlass/epilogue/block/block_epilogue_mla_fd_rescale_o.hpp+7-7
@@ -46,9 +46,9 @@ public:
46 static constexpr uint32_t STAGES = 2;46 static constexpr uint32_t STAGES = 2;
47 47 
48 CATLASS_DEVICE48 CATLASS_DEVICE
49- BlockEpilogue(Arch::Resource<ArchTag> &resource, uint32_t kvSplitCoreNum_)49+ BlockEpilogue(Arch::Resource<ArchTag> &resource, uint32_t maxKvSplitCoreNum_)
50 {50 {
51- kvSplitCoreNum = kvSplitCoreNum_;51+ maxKvSplitCoreNum = maxKvSplitCoreNum_;
52 52 
53 uint32_t ubOffset = 0;53 uint32_t ubOffset = 0;
54 oIn[0] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);54 oIn[0] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
@@ -112,7 +112,7 @@ public:
112 AscendC::GlobalTensor<ElementOutput> gOutput,112 AscendC::GlobalTensor<ElementOutput> gOutput,
113 AscendC::GlobalTensor<ElementInput> gOCoreTmp,113 AscendC::GlobalTensor<ElementInput> gOCoreTmp,
114 AscendC::GlobalTensor<ElementInput> gl,114 AscendC::GlobalTensor<ElementInput> gl,
115- uint32_t actualHeads, uint32_t headsProcess, uint32_t headSize)115+ uint32_t actualHeads, uint32_t headsProcess, uint32_t headSize, uint32_t kvSplitCoreNum)
116 {116 {
117 uint32_t kvSplitRound = (kvSplitCoreNum + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE * FLOAT_BLOCK_SIZE;117 uint32_t kvSplitRound = (kvSplitCoreNum + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE * FLOAT_BLOCK_SIZE;
118 118 
@@ -120,7 +120,7 @@ public:
120 AscendC::DataCopyPad(120 AscendC::DataCopyPad(
121 lIn, gl,121 lIn, gl,
122 AscendC::DataCopyExtParams(122 AscendC::DataCopyExtParams(
123- actualHeads, kvSplitCoreNum * sizeof(ElementInput), 0,123+ actualHeads, kvSplitCoreNum * sizeof(ElementInput), (maxKvSplitCoreNum - kvSplitCoreNum) * sizeof(ElementInput),
124 (KV_SPLIT_MAX - kvSplitCoreNum) / FLOAT_BLOCK_SIZE, 0),124 (KV_SPLIT_MAX - kvSplitCoreNum) / FLOAT_BLOCK_SIZE, 0),
125 AscendC::DataCopyPadExtParams<ElementInput>(false, 0, 0, 0));125 AscendC::DataCopyPadExtParams<ElementInput>(false, 0, 0, 0));
126 126 
@@ -205,7 +205,7 @@ public:
205 oIn[0], gOCoreTmp,205 oIn[0], gOCoreTmp,
206 AscendC::DataCopyExtParams(206 AscendC::DataCopyExtParams(
207 actualHeads, headSize * sizeof(ElementInput),207 actualHeads, headSize * sizeof(ElementInput),
208- (kvSplitCoreNum * headSize - headSize) * sizeof(ElementInput), 0, 0),208+ (maxKvSplitCoreNum * headSize - headSize) * sizeof(ElementInput), 0, 0),
209 AscendC::DataCopyPadExtParams<ElementInput>(false, 0, 0, 0));209 AscendC::DataCopyPadExtParams<ElementInput>(false, 0, 0, 0));
210 AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);210 AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
211 211 
@@ -223,7 +223,7 @@ public:
223 oIn[nextBufferId], gOCoreTmp[(i + 1) * headSize],223 oIn[nextBufferId], gOCoreTmp[(i + 1) * headSize],
224 AscendC::DataCopyExtParams(224 AscendC::DataCopyExtParams(
225 actualHeads, headSize * sizeof(ElementInput),225 actualHeads, headSize * sizeof(ElementInput),
226- (kvSplitCoreNum * headSize - headSize) * sizeof(ElementInput), 0, 0),226+ (maxKvSplitCoreNum * headSize - headSize) * sizeof(ElementInput), 0, 0),
227 AscendC::DataCopyPadExtParams<ElementInput>(false, 0, 0, 0));227 AscendC::DataCopyPadExtParams<ElementInput>(false, 0, 0, 0));
228 AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(oInEventList[nextBufferId]);228 AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(oInEventList[nextBufferId]);
229 }229 }
@@ -289,7 +289,7 @@ public:
289 }289 }
290 290 
291private:291private:
292- uint32_t kvSplitCoreNum = 1;292+ uint32_t maxKvSplitCoreNum = 1;
293 AscendC::LocalTensor<ElementOutput> out;293 AscendC::LocalTensor<ElementOutput> out;
294 AscendC::LocalTensor<float> oIn[STAGES];294 AscendC::LocalTensor<float> oIn[STAGES];
295 AscendC::LocalTensor<float> oTemp[STAGES];295 AscendC::LocalTensor<float> oTemp[STAGES];
Minclude/catlass/epilogue/block/block_epilogue_mla_rescale_o.hpp+79-17
@@ -80,7 +80,8 @@ public:
80 goUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(GO_UB_TENSOR_OFFSET);80 goUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(GO_UB_TENSOR_OFFSET);
81 goUbTensor16 = resource.ubBuf.template GetBufferByByte<ElementOutput>(GO_UB_TENSOR_OFFSET);81 goUbTensor16 = resource.ubBuf.template GetBufferByByte<ElementOutput>(GO_UB_TENSOR_OFFSET);
82 hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);82 hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);
83- gmUbTensor = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);83+ gmUbTensor[0] = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
84+ gmUbTensor[1] = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET + UB_UINT8_LINE_SIZE);
84 }85 }
85 86 
86 CATLASS_DEVICE87 CATLASS_DEVICE
@@ -125,7 +126,7 @@ public:
125 uint32_t epiTokenNum,126 uint32_t epiTokenNum,
126 uint32_t integralHeadNum,127 uint32_t integralHeadNum,
127 uint32_t rescaleOPingPongFlag,128 uint32_t rescaleOPingPongFlag,
128- uint32_t &glFlag)129+ uint32_t *glFlag, uint32_t taskPingPongFlag)
129 {130 {
130 uint32_t curRowNum = layoutInput.shape(0);131 uint32_t curRowNum = layoutInput.shape(0);
131 uint32_t embed = layoutInput.shape(1);132 uint32_t embed = layoutInput.shape(1);
@@ -139,6 +140,8 @@ public:
139 uint64_t llUbOffsetCurCycle = (uint64_t)(rescaleOPingPongFlag * HALF_LL_UB_SIZE +140 uint64_t llUbOffsetCurCycle = (uint64_t)(rescaleOPingPongFlag * HALF_LL_UB_SIZE +
140 rowLoopIdx * ROW_WISE_CYCLE_TILE);141 rowLoopIdx * ROW_WISE_CYCLE_TILE);
141 uint32_t oUbOffset = oPingPangFlag * ROW_WISE_CYCLE_TILE * embedRound;142 uint32_t oUbOffset = oPingPangFlag * ROW_WISE_CYCLE_TILE * embedRound;
143+ uint32_t qHeads = strideQO / embed;
144+ 
142 AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(oPingPangFlag);145 AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(oPingPangFlag);
143 if ((nIdx - 1) != 0) {146 if ((nIdx - 1) != 0) {
144 AscendC::DataCopy(147 AscendC::DataCopy(
@@ -298,7 +301,7 @@ public:
298 AscendC::PipeBarrier<PIPE_V>();301 AscendC::PipeBarrier<PIPE_V>();
299 AscendC::Brcb(302 AscendC::Brcb(
300 hmUbTensor.ReinterpretCast<uint32_t>(),303 hmUbTensor.ReinterpretCast<uint32_t>(),
301- gmUbTensor.ReinterpretCast<uint32_t>()[rowLoopIdx * ROW_WISE_CYCLE_TILE],304+ gmUbTensor[taskPingPongFlag].ReinterpretCast<uint32_t>()[rowLoopIdx * ROW_WISE_CYCLE_TILE],
302 curRowNumRound / FLOAT_BLOCK_SIZE,305 curRowNumRound / FLOAT_BLOCK_SIZE,
303 AscendC::BrcbRepeatParams(1, 8));306 AscendC::BrcbRepeatParams(1, 8));
304 AscendC::PipeBarrier<PIPE_V>();307 AscendC::PipeBarrier<PIPE_V>();
@@ -314,18 +317,76 @@ public:
314 317 
315 AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID2);318 AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID2);
316 AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID2);319 AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID2);
317- AscendC::DataCopyPad(gl, tvUbTensor,
318- AscendC::DataCopyExtParams(curRowNum, 4, 0, (kvSplitCoreNum - 1) * 4, 0));
319 320 
320- if (glFlag == 0) {321+ if (tokenNumPerHead == 1) {
321- AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID2);322+ AscendC::DataCopyPad(gl, tvUbTensor,
322- glFlag = 1;323+ AscendC::DataCopyExtParams(curRowNum, 4, 0, (kvSplitCoreNum - 1) * 4, 0));
324+ } else {
325+ uint32_t innerOGmOffset = 0;
326+ uint32_t inner_go_ubuf_offset = 0;
327+ if (proTokenNum != 0) {
328+ AscendC::DataCopyPad(gl[innerOGmOffset + proTokenIdx * kvSplitCoreNum * qHeads],
329+ tvUbTensor[inner_go_ubuf_offset],
330+ AscendC::DataCopyExtParams(proTokenNum,
331+ 4, 0, (kvSplitCoreNum * qHeads - 1) * 4, 0));
332+ innerOGmOffset += kvSplitCoreNum;
333+ inner_go_ubuf_offset += proTokenNum * FLOAT_BLOCK_SIZE;
334+ }
335+ for (uint32_t qN_idx = 0; qN_idx < integralHeadNum; qN_idx++) {
336+ AscendC::DataCopyPad(gl[innerOGmOffset],
337+ tvUbTensor[inner_go_ubuf_offset],
338+ AscendC::DataCopyExtParams(tokenNumPerHead,
339+ 4, 0, (kvSplitCoreNum * qHeads - 1) * 4, 0));
340+ innerOGmOffset += kvSplitCoreNum;
341+ inner_go_ubuf_offset += tokenNumPerHead * FLOAT_BLOCK_SIZE;
342+ }
343+ 
344+ if (epiTokenNum != 0) {
345+ AscendC::DataCopyPad(gl[innerOGmOffset],
346+ tvUbTensor[inner_go_ubuf_offset],
347+ AscendC::DataCopyExtParams(epiTokenNum,
348+ 4, 0, (kvSplitCoreNum * qHeads - 1) * 4, 0));
349+ }
350+ }
351+ 
352+ if (glFlag[taskPingPongFlag] == 0) {
353+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(taskPingPongFlag + 2);
354+ glFlag[taskPingPongFlag] = 1;
323 }355 }
324 uint32_t srcGap = ((embed % 16 <= 8) && (embed % 16 > 0)) ? 1 : 0;356 uint32_t srcGap = ((embed % 16 <= 8) && (embed % 16 > 0)) ? 1 : 0;
325- AscendC::DataCopyPad(gOCoreTmp, goUbTensor32[oUbOffset],357+ 
326- AscendC::DataCopyExtParams(curRowNum, embed * 4, srcGap, (kvSplitCoreNum - 1) * embed * 4, 0));358+ if (tokenNumPerHead == 1) {
327- AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID3);359+ AscendC::DataCopyPad(gOCoreTmp, goUbTensor32[oUbOffset],
328- AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID3);360+ AscendC::DataCopyExtParams(curRowNum, embed * 4, srcGap, (kvSplitCoreNum - 1) * embed * 4, 0));
361+ } else {
362+ uint32_t innerOGmOffset = 0;
363+ uint32_t inner_go_ubuf_offset = oUbOffset;
364+ if (proTokenNum != 0) {
365+ AscendC::DataCopyPad(gOCoreTmp[innerOGmOffset + proTokenIdx * kvSplitCoreNum * strideQO],
366+ goUbTensor32[inner_go_ubuf_offset],
367+ AscendC::DataCopyExtParams(proTokenNum,
368+ embed * 4, 0, (kvSplitCoreNum * strideQO - embed) * 4, 0));
369+ innerOGmOffset += embed * kvSplitCoreNum;
370+ inner_go_ubuf_offset += proTokenNum * embed;
371+ }
372+ for (uint32_t qN_idx = 0; qN_idx < integralHeadNum; qN_idx++) {
373+ AscendC::DataCopyPad(gOCoreTmp[innerOGmOffset],
374+ goUbTensor32[inner_go_ubuf_offset],
375+ AscendC::DataCopyExtParams(tokenNumPerHead,
376+ embed * 4, 0, (kvSplitCoreNum * strideQO - embed) * 4, 0));
377+ innerOGmOffset += embed * kvSplitCoreNum;
378+ inner_go_ubuf_offset += tokenNumPerHead * embed;
379+ }
380+ 
381+ if (epiTokenNum != 0) {
382+ AscendC::DataCopyPad(gOCoreTmp[innerOGmOffset],
383+ goUbTensor32[inner_go_ubuf_offset],
384+ AscendC::DataCopyExtParams(epiTokenNum,
385+ embed * 4, 0, (kvSplitCoreNum * strideQO - embed) * 4, 0));
386+ }
387+ }
388+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
389+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
329 AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID1);390 AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID1);
330 AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID1);391 AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID1);
331 } else {392 } else {
@@ -406,7 +467,7 @@ public:
406 uint32_t isLastNTile,467 uint32_t isLastNTile,
407 uint32_t curHeadNum,468 uint32_t curHeadNum,
408 uint32_t rescaleOPingPongFlag,469 uint32_t rescaleOPingPongFlag,
409- uint32_t &glFlag)470+ uint32_t *glFlag, uint32_t taskPingPongFlag)
410 {471 {
411 uint32_t tokenNumPerHead = layoutOutput.shape(0);472 uint32_t tokenNumPerHead = layoutOutput.shape(0);
412 uint32_t embed = layoutInput.shape(1);473 uint32_t embed = layoutInput.shape(1);
@@ -450,12 +511,13 @@ public:
450 proTokenNum = (tokenNumPerHead - epiTokenNum) % tokenNumPerHead;511 proTokenNum = (tokenNumPerHead - epiTokenNum) % tokenNumPerHead;
451 integralHeadNum = (rowActualCurCycle - proTokenNum) / tokenNumPerHead;512 integralHeadNum = (rowActualCurCycle - proTokenNum) / tokenNumPerHead;
452 epiTokenNum = rowActualCurCycle - proTokenNum - integralHeadNum * tokenNumPerHead;513 epiTokenNum = rowActualCurCycle - proTokenNum - integralHeadNum * tokenNumPerHead;
514+ int64_t headIdx = rowOffsetLoop / tokenNumPerHead;
453 SubCoreCompute(gInputThisCurCycle, gUpdateCurCycle, gOutputCurCycle,515 SubCoreCompute(gInputThisCurCycle, gUpdateCurCycle, gOutputCurCycle,
454- gOCoreTmp[rowOffsetLoop * embed * kvSplitCoreNum],516+ gOCoreTmp[headIdx * kvSplitCoreNum * embed],
455- gl[rowOffsetLoop * kvSplitCoreNum],517+ gl[headIdx * kvSplitCoreNum],
456 layoutInputCurCycle, layoutOutputCurCycle, layoutUpdateCurCycle,518 layoutInputCurCycle, layoutOutputCurCycle, layoutUpdateCurCycle,
457 nIdx, isLastNTile, needRowLoop, rowLoopIdx,519 nIdx, isLastNTile, needRowLoop, rowLoopIdx,
458- proTokenIdx, proTokenNum, epiTokenNum, integralHeadNum, rescaleOPingPongFlag, glFlag);520+ proTokenIdx, proTokenNum, epiTokenNum, integralHeadNum, rescaleOPingPongFlag, glFlag, taskPingPongFlag);
459 }521 }
460 }522 }
461 }523 }
@@ -471,7 +533,7 @@ private:
471 AscendC::LocalTensor<float> tvUbTensor;533 AscendC::LocalTensor<float> tvUbTensor;
472 AscendC::LocalTensor<float> goUbTensor32;534 AscendC::LocalTensor<float> goUbTensor32;
473 AscendC::LocalTensor<float> hmUbTensor;535 AscendC::LocalTensor<float> hmUbTensor;
474- AscendC::LocalTensor<float> gmUbTensor;536+ AscendC::LocalTensor<float> gmUbTensor[2];
475};537};
476 538 
477} // namespace Catlass::Epilogue::Block539} // namespace Catlass::Epilogue::Block
Minclude/catlass/epilogue/block/block_epilogue_mla_softmax.hpp+119-49
@@ -64,13 +64,15 @@ public:
64 static constexpr uint32_t HALF_DM_UB_SIZE = 128;64 static constexpr uint32_t HALF_DM_UB_SIZE = 128;
65 static constexpr uint32_t VECTOR_SIZE = 128;65 static constexpr uint32_t VECTOR_SIZE = 128;
66 static constexpr uint32_t HALF_LL_UB_SIZE = 256;66 static constexpr uint32_t HALF_LL_UB_SIZE = 256;
67+ static constexpr uint32_t REDUCE_UB_SIZE = 1024;
67 68 
68 CATLASS_DEVICE69 CATLASS_DEVICE
69 BlockEpilogue(Arch::Resource<ArchTag> &resource, half tor_, uint32_t kvSplitCoreNum_)70 BlockEpilogue(Arch::Resource<ArchTag> &resource, half tor_, uint32_t kvSplitCoreNum_)
70 {71 {
71 // Allocate UB space72 // Allocate UB space
72 constexpr uint32_t LS_UB_TENSOR_OFFSET = 0;73 constexpr uint32_t LS_UB_TENSOR_OFFSET = 0;
73- constexpr uint32_t LP_UB_TENSOR_OFFSET = 2 * UB_UINT8_BLOCK_SIZE_MLA;74+ constexpr uint32_t LP_UB_TENSOR_OFFSET = 0;
75+ // constexpr uint32_t LP_UB_TENSOR_OFFSET = 2 * UB_UINT8_BLOCK_SIZE_MLA;
74 constexpr uint32_t LM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA;76 constexpr uint32_t LM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA;
75 constexpr uint32_t HM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE;77 constexpr uint32_t HM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE;
76 constexpr uint32_t DM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 6 * UB_UINT8_LINE_SIZE;78 constexpr uint32_t DM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 6 * UB_UINT8_LINE_SIZE;
@@ -85,7 +87,8 @@ public:
85 lsUbTensor = resource.ubBuf.template GetBufferByByte<float>(LS_UB_TENSOR_OFFSET);87 lsUbTensor = resource.ubBuf.template GetBufferByByte<float>(LS_UB_TENSOR_OFFSET);
86 lmUbTensor = resource.ubBuf.template GetBufferByByte<float>(LM_UB_TENSOR_OFFSET);88 lmUbTensor = resource.ubBuf.template GetBufferByByte<float>(LM_UB_TENSOR_OFFSET);
87 hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);89 hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);
88- gmUbTensor = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);90+ gmUbTensor[0] = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
91+ gmUbTensor[1] = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET + UB_UINT8_LINE_SIZE);
89 dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);92 dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
90 llUbTensor = resource.ubBuf.template GetBufferByByte<float>(LL_UB_TENSOR_OFFSET);93 llUbTensor = resource.ubBuf.template GetBufferByByte<float>(LL_UB_TENSOR_OFFSET);
91 tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);94 tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
@@ -96,6 +99,38 @@ public:
96 {99 {
97 }100 }
98 101 
102+ CATLASS_DEVICE
103+ void SetVecMask(int32_t len)
104+ {
105+ uint64_t mask = 0;
106+ uint64_t one = 1;
107+ uint64_t temp = len % FLOAT_VECTOR_SIZE;
108+ for (int64_t i = 0; i < temp; i++) {
109+ mask |= one << i;
110+ }
111+ 
112+ if (len == VECTOR_SIZE || len == 0) {
113+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
114+ } else if (len >= FLOAT_VECTOR_SIZE) {
115+ AscendC::SetVectorMask<int8_t>(mask, (uint64_t)-1);
116+ } else {
117+ AscendC::SetVectorMask<int8_t>(0x0, mask);
118+ }
119+ }
120+ 
121+ CATLASS_DEVICE
122+ void SetBlockReduceMask(int32_t len)
123+ {
124+ if (len > 8 || len < 1) {
125+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
126+ return;
127+ }
128+ uint64_t subMask = ((uint64_t)1 << len) - 1;
129+ uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask + (subMask << 56) +
130+ (subMask << 40) + (subMask << 24) + (subMask << 8);
131+ AscendC::SetVectorMask<int8_t>(maskValue, maskValue);
132+ }
133+ 
99 CATLASS_DEVICE134 CATLASS_DEVICE
100 void SetMask(int32_t len)135 void SetMask(int32_t len)
101 {136 {
@@ -196,6 +231,59 @@ public:
196 AscendC::PipeBarrier<PIPE_V>();231 AscendC::PipeBarrier<PIPE_V>();
197 }232 }
198 233 
234+ CATLASS_DEVICE
235+ void RowsumTAILTILE(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowsumUb,
236+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
237+ uint32_t numElemsAligned)
238+ {
239+ if (numElems >= FLOAT_VECTOR_SIZE) {
240+ AscendC::BlockReduceSum<float, false>(tvUbTensor, srcUb, numRowsRound, 0, 1, 1,
241+ numElemsAligned / FLOAT_BLOCK_SIZE);
242+ AscendC::PipeBarrier<PIPE_V>();
243+ AscendC::BlockReduceSum<float, false>(
244+ rowsumUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0,
245+ 1, 1, 8);
246+ AscendC::PipeBarrier<PIPE_V>();
247+ for (uint64_t rowSumIdx = 1; rowSumIdx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowSumIdx) {
248+ AscendC::BlockReduceSum<float, false>(tvUbTensor, srcUb[rowSumIdx * FLOAT_VECTOR_SIZE], numRowsRound, 0,
249+ 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
250+ AscendC::PipeBarrier<PIPE_V>();
251+ AscendC::BlockReduceSum<float, false>(
252+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
253+ (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
254+ AscendC::PipeBarrier<PIPE_V>();
255+ SetVecMask(numRowsRound);
256+ AscendC::Add<float, false>(rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
257+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
258+ AscendC::PipeBarrier<PIPE_V>();
259+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
260+ }
261+ }
262+ if (numElems % FLOAT_VECTOR_SIZE > 0) {
263+ SetVecMask(numElems % FLOAT_VECTOR_SIZE);
264+ AscendC::BlockReduceSum<float, false>(tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
265+ numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
266+ AscendC::PipeBarrier<PIPE_V>();
267+ SetBlockReduceMask((numElems % FLOAT_VECTOR_SIZE + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE);
268+ if (numElems < FLOAT_VECTOR_SIZE) {
269+ AscendC::BlockReduceSum<float, false>(
270+ rowsumUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
271+ 0, 1, 1, 8);
272+ AscendC::PipeBarrier<PIPE_V>();
273+ } else {
274+ AscendC::BlockReduceSum<float, false>(
275+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
276+ (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
277+ AscendC::PipeBarrier<PIPE_V>();
278+ SetVecMask(numRowsRound);
279+ AscendC::Add<float, false>(rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
280+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
281+ AscendC::PipeBarrier<PIPE_V>();
282+ }
283+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
284+ }
285+ }
286+ 
199 CATLASS_DEVICE287 CATLASS_DEVICE
200 void ReduceMaxRepeatM(288 void ReduceMaxRepeatM(
201 const AscendC::LocalTensor<float> &dst,289 const AscendC::LocalTensor<float> &dst,
@@ -253,7 +341,7 @@ public:
253 const LayoutInput &layoutInput,341 const LayoutInput &layoutInput,
254 uint32_t nIdx,342 uint32_t nIdx,
255 uint32_t softmaxPingPongFlag,343 uint32_t softmaxPingPongFlag,
256- uint32_t &glFlag)344+ uint32_t *glFlag, uint32_t taskPingPongFlag, uint32_t sUbOffset)
257 {345 {
258 uint32_t curRowNum = layoutInput.shape(0);346 uint32_t curRowNum = layoutInput.shape(0);
259 uint32_t kSeqTile = layoutInput.shape(1);347 uint32_t kSeqTile = layoutInput.shape(1);
@@ -262,45 +350,27 @@ public:
262 uint32_t sub_m_d64 = (curRowNum + 63) / 64; // up aligned to 128350 uint32_t sub_m_d64 = (curRowNum + 63) / 64; // up aligned to 128
263 uint64_t dmUbOffsetCurCycle = (uint64_t)(softmaxPingPongFlag * HALF_DM_UB_SIZE);351 uint64_t dmUbOffsetCurCycle = (uint64_t)(softmaxPingPongFlag * HALF_DM_UB_SIZE);
264 uint64_t llUbOffsetCurCycle = (uint64_t)(softmaxPingPongFlag * HALF_LL_UB_SIZE);352 uint64_t llUbOffsetCurCycle = (uint64_t)(softmaxPingPongFlag * HALF_LL_UB_SIZE);
265- AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);353+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(softmaxPingPongFlag + 6);
266- AscendC::DataCopy(lsUbTensor, gInput,354+ AscendC::DataCopy(lsUbTensor[sUbOffset], gInput,
267 AscendC::DataCopyParams(1, curRowNum * kSeqTileRound / FLOAT_BLOCK_SIZE, 0, 0));355 AscendC::DataCopyParams(1, curRowNum * kSeqTileRound / FLOAT_BLOCK_SIZE, 0, 0));
268 356 
269 AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);357 AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
270 AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);358 AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
271 359 
272 // muls scale_value360 // muls scale_value
273- for (uint32_t mulsIdx = 0; mulsIdx < kSeqTile / FLOAT_VECTOR_SIZE; ++mulsIdx) {361+ AscendC::Muls<float, false>(lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], tor, (uint64_t)0,
274- AscendC::Muls<float, false>(362+ (curRowNum * kSeqTileRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
275- lsUbTensor[mulsIdx * FLOAT_VECTOR_SIZE],363+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
276- lsUbTensor[mulsIdx * FLOAT_VECTOR_SIZE],
277- tor,
278- (uint64_t)0,
279- curRowNum,
280- AscendC::UnaryRepeatParams(1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE));
281- }
282- if (kSeqTile % FLOAT_VECTOR_SIZE > 0) {
283- SetMask(kSeqTile % FLOAT_VECTOR_SIZE);
284- AscendC::Muls<float, false>(
285- lsUbTensor[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
286- lsUbTensor[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
287- tor,
288- (uint64_t)0,
289- curRowNum,
290- AscendC::UnaryRepeatParams(1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE));
291- 
292- AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
293- }
294 AscendC::PipeBarrier<PIPE_V>();364 AscendC::PipeBarrier<PIPE_V>();
295 365 
296 // *** lm = rowmax(ls)366 // *** lm = rowmax(ls)
297- ReduceMaxRepeatM(lmUbTensor, lsUbTensor, lpUbTensor32, curRowNum, kSeqTile, kSeqTileRound);367+ ReduceMaxRepeatM(lmUbTensor, lsUbTensor[sUbOffset], tvUbTensor, curRowNum, kSeqTile, kSeqTileRound);
298 368 
299 if (nIdx != 0) {369 if (nIdx != 0) {
300 AscendC::Max<float, false>(370 AscendC::Max<float, false>(
301 hmUbTensor,371 hmUbTensor,
302 lmUbTensor,372 lmUbTensor,
303- gmUbTensor,373+ gmUbTensor[taskPingPongFlag],
304 (uint64_t)0,374 (uint64_t)0,
305 sub_m_d64,375 sub_m_d64,
306 AscendC::BinaryRepeatParams(376 AscendC::BinaryRepeatParams(
@@ -308,7 +378,7 @@ public:
308 AscendC::PipeBarrier<PIPE_V>();378 AscendC::PipeBarrier<PIPE_V>();
309 AscendC::Sub<float, false>(379 AscendC::Sub<float, false>(
310 dmUbTensor[dmUbOffsetCurCycle],380 dmUbTensor[dmUbOffsetCurCycle],
311- gmUbTensor,381+ gmUbTensor[taskPingPongFlag],
312 hmUbTensor,382 hmUbTensor,
313 (uint64_t)0,383 (uint64_t)0,
314 sub_m_d64,384 sub_m_d64,
@@ -320,44 +390,45 @@ public:
320 AscendC::PipeBarrier<PIPE_V>();390 AscendC::PipeBarrier<PIPE_V>();
321 }391 }
322 // *** gm = hm392 // *** gm = hm
323- AscendC::DataCopy(gmUbTensor, hmUbTensor, AscendC::DataCopyParams(1, subMRound / FLOAT_BLOCK_SIZE, 0, 0));393+ AscendC::DataCopy(gmUbTensor[taskPingPongFlag], hmUbTensor, AscendC::DataCopyParams(1, subMRound / FLOAT_BLOCK_SIZE, 0, 0));
324 AscendC::PipeBarrier<PIPE_V>();394 AscendC::PipeBarrier<PIPE_V>();
325 395 
326 if (kvSplitCoreNum != 1) {396 if (kvSplitCoreNum != 1) {
327 if (nIdx == 0) {397 if (nIdx == 0) {
328- if (glFlag == 1) {398+ if (glFlag[taskPingPongFlag] == 1) {
329- AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID2);399+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(taskPingPongFlag + 2);
330- glFlag = 0;400+ glFlag[taskPingPongFlag] = 0;
331 }401 }
332 }402 }
333 }403 }
334 404 
335 // *** ls = ls - hm_block405 // *** ls = ls - hm_block
336- TensorSubValueRepeatM(lsUbTensor, lsUbTensor,406+ TensorSubValueRepeatM(lsUbTensor[sUbOffset], lsUbTensor[sUbOffset],
337 hmUbTensor, tvUbTensor,407 hmUbTensor, tvUbTensor,
338 curRowNum, subMRound, kSeqTile, kSeqTileRound);408 curRowNum, subMRound, kSeqTile, kSeqTileRound);
339 409 
340 AscendC::Exp<float, false>(410 AscendC::Exp<float, false>(
341- lsUbTensor,411+ lsUbTensor[sUbOffset],
342- lsUbTensor,412+ lsUbTensor[sUbOffset],
343 (uint64_t)0,413 (uint64_t)0,
344 (curRowNum * kSeqTileRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,414 (curRowNum * kSeqTileRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
345 AscendC::UnaryRepeatParams(1, 1, 8, 8));415 AscendC::UnaryRepeatParams(1, 1, 8, 8));
346 416 
347 AscendC::PipeBarrier<PIPE_V>();417 AscendC::PipeBarrier<PIPE_V>();
418+ RowsumTAILTILE(lsUbTensor[sUbOffset], llUbTensor[llUbOffsetCurCycle], tvUbTensor, curRowNum, kSeqTile, kSeqTileRound);
419+ 
348 // *** lp = castfp32to16(ls)420 // *** lp = castfp32to16(ls)
349 if (std::is_same<ElementOutput, bfloat16_t>::value) {421 if (std::is_same<ElementOutput, bfloat16_t>::value) {
350 AscendC::Cast<ElementOutput, float, false>(422 AscendC::Cast<ElementOutput, float, false>(
351- tvUbTensor16, lsUbTensor, AscendC::RoundMode::CAST_RINT, (uint64_t)0,423+ tvUbTensor16[sUbOffset * 2], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_RINT, (uint64_t)0,
352 (curRowNum * kSeqTileRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,424 (curRowNum * kSeqTileRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
353 AscendC::UnaryRepeatParams(1, 1, 4, 8));425 AscendC::UnaryRepeatParams(1, 1, 4, 8));
354 } else {426 } else {
355 AscendC::Cast<ElementOutput, float, false>(427 AscendC::Cast<ElementOutput, float, false>(
356- tvUbTensor16, lsUbTensor, AscendC::RoundMode::CAST_NONE, (uint64_t)0,428+ tvUbTensor16[sUbOffset * 2], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_NONE, (uint64_t)0,
357 (curRowNum * kSeqTileRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,429 (curRowNum * kSeqTileRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
358 AscendC::UnaryRepeatParams(1, 1, 4, 8));430 AscendC::UnaryRepeatParams(1, 1, 4, 8));
359 }431 }
360- 
361 AscendC::PipeBarrier<PIPE_V>();432 AscendC::PipeBarrier<PIPE_V>();
362 433 
363 AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);434 AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
@@ -369,17 +440,15 @@ public:
369 uint16_t dstStride = 0;440 uint16_t dstStride = 0;
370 441 
371 AscendC::DataCopy(gOutput,442 AscendC::DataCopy(gOutput,
372- tvUbTensor16,443+ tvUbTensor16[sUbOffset * 2],
373 AscendC::DataCopyParams(444 AscendC::DataCopyParams(
374- blockCount, // blockCount445+ blockCount,
375- blockLen, // blockLen446+ blockLen,
376- srcStride, // srcGap447+ srcStride,
377 dstStride));448 dstStride));
378 449 
379 // *** ll = rowsum(ls32)450 // *** ll = rowsum(ls32)
380- ReduceSumRepeatM(llUbTensor[llUbOffsetCurCycle], lsUbTensor, curRowNum, kSeqTile, kSeqTileRound);451+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(softmaxPingPongFlag + 6);
381- AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);
382- AscendC::PipeBarrier<PIPE_V>();
383 }452 }
384 453 
385 CATLASS_DEVICE454 CATLASS_DEVICE
@@ -392,7 +461,7 @@ public:
392 uint32_t nIdx,461 uint32_t nIdx,
393 uint32_t curHeadNum,462 uint32_t curHeadNum,
394 uint32_t softmaxPingPongFlag,463 uint32_t softmaxPingPongFlag,
395- uint32_t &glFlag)464+ uint32_t *glFlag, uint32_t taskPingPongFlag)
396 {465 {
397 uint32_t rowActual = actualBlockShape.m();466 uint32_t rowActual = actualBlockShape.m();
398 uint32_t nActual = actualBlockShape.n();467 uint32_t nActual = actualBlockShape.n();
@@ -400,6 +469,7 @@ public:
400 469 
401 uint32_t subBlockIdx = AscendC::GetSubBlockIdx();470 uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
402 uint32_t subBlockNum = AscendC::GetSubBlockNum();471 uint32_t subBlockNum = AscendC::GetSubBlockNum();
472+ uint32_t sUbOffset = softmaxPingPongFlag * 8192;
403 473 
404 uint32_t curHeadSplitSubBlock = curHeadNum / subBlockNum;474 uint32_t curHeadSplitSubBlock = curHeadNum / subBlockNum;
405 uint32_t curHeadThisSubBlock = (subBlockIdx == 0) ? curHeadSplitSubBlock : (curHeadNum - curHeadSplitSubBlock);475 uint32_t curHeadThisSubBlock = (subBlockIdx == 0) ? curHeadSplitSubBlock : (curHeadNum - curHeadSplitSubBlock);
@@ -415,7 +485,7 @@ public:
415 auto gOutputThisSubBlock = gOutput[offsetOutput];485 auto gOutputThisSubBlock = gOutput[offsetOutput];
416 auto layoutOutputThisSubBlock = layoutOutput.GetTileLayout(MatrixCoord(rowActualThisSubBlock, nActual));486 auto layoutOutputThisSubBlock = layoutOutput.GetTileLayout(MatrixCoord(rowActualThisSubBlock, nActual));
417 SubCoreCompute(gOutputThisSubBlock, gInputThisSubBlock, layoutOutputThisSubBlock, layoutInputThisSubBlock,487 SubCoreCompute(gOutputThisSubBlock, gInputThisSubBlock, layoutOutputThisSubBlock, layoutInputThisSubBlock,
418- nIdx, softmaxPingPongFlag, glFlag);488+ nIdx, softmaxPingPongFlag, glFlag, taskPingPongFlag, sUbOffset);
419 }489 }
420 }490 }
421 491 
@@ -428,7 +498,7 @@ private:
428 AscendC::LocalTensor<float> lsUbTensor;498 AscendC::LocalTensor<float> lsUbTensor;
429 AscendC::LocalTensor<float> lmUbTensor;499 AscendC::LocalTensor<float> lmUbTensor;
430 AscendC::LocalTensor<float> hmUbTensor;500 AscendC::LocalTensor<float> hmUbTensor;
431- AscendC::LocalTensor<float> gmUbTensor;501+ AscendC::LocalTensor<float> gmUbTensor[2];
432 AscendC::LocalTensor<float> dmUbTensor;502 AscendC::LocalTensor<float> dmUbTensor;
433 AscendC::LocalTensor<float> llUbTensor;503 AscendC::LocalTensor<float> llUbTensor;
434 AscendC::LocalTensor<float> tvUbTensor;504 AscendC::LocalTensor<float> tvUbTensor;
Minclude/catlass/epilogue/block/block_epilogue_mla_tp1_rescale_o.hpp+30-57
@@ -1,40 +1,33 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.8+ * See LICENSE in the root of the software repository for the full text of the License.
9- */9+ */
10- 
11#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_TP1_RESCALE_O_HPP10#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_TP1_RESCALE_O_HPP
12#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_TP1_RESCALE_O_HPP11#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_TP1_RESCALE_O_HPP
13- 
14#include "catlass/catlass.hpp"12#include "catlass/catlass.hpp"
15#include "catlass/arch/resource.hpp"13#include "catlass/arch/resource.hpp"
16#include "catlass/epilogue/dispatch_policy.hpp"14#include "catlass/epilogue/dispatch_policy.hpp"
17#include "catlass/epilogue/tile/tile_copy.hpp"15#include "catlass/epilogue/tile/tile_copy.hpp"
18#include "catlass/gemm_coord.hpp"16#include "catlass/gemm_coord.hpp"
19#include "catlass/matrix_coord.hpp"17#include "catlass/matrix_coord.hpp"
20- 
21namespace Catlass::Epilogue::Block {18namespace Catlass::Epilogue::Block {
22- 
23template <class OutputType_, class UpdateType_, class InputType_>19template <class OutputType_, class UpdateType_, class InputType_>
24class BlockEpilogue<EpilogueAtlasA2MLATP1RescaleO, OutputType_, UpdateType_, InputType_> {20class BlockEpilogue<EpilogueAtlasA2MLATP1RescaleO, OutputType_, UpdateType_, InputType_> {
25public:21public:
26 // Type aliases22 // Type aliases
27 using DispatchPolicy = EpilogueAtlasA2MLATP1RescaleO;23 using DispatchPolicy = EpilogueAtlasA2MLATP1RescaleO;
28 using ArchTag = typename DispatchPolicy::ArchTag;24 using ArchTag = typename DispatchPolicy::ArchTag;
29- 
30 using ElementOutput = typename OutputType_::Element;25 using ElementOutput = typename OutputType_::Element;
31 using ElementUpdate = typename UpdateType_::Element;26 using ElementUpdate = typename UpdateType_::Element;
32 using ElementInput = typename InputType_::Element;27 using ElementInput = typename InputType_::Element;
33- 
34 using LayoutOutput = typename OutputType_::Layout;28 using LayoutOutput = typename OutputType_::Layout;
35 using LayoutUpdate = typename UpdateType_::Layout;29 using LayoutUpdate = typename UpdateType_::Layout;
36 using LayoutInput = typename InputType_::Layout;30 using LayoutInput = typename InputType_::Layout;
37- 
38 static constexpr uint32_t HALF_ELENUM_PER_BLK = 16;31 static constexpr uint32_t HALF_ELENUM_PER_BLK = 16;
39 static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128;32 static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128;
40 static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64;33 static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64;
@@ -50,33 +43,31 @@ public:
50 static constexpr uint32_t HALF_LL_UB_SIZE = 256;43 static constexpr uint32_t HALF_LL_UB_SIZE = 256;
51 static constexpr uint32_t VECTOR_SIZE = 128;44 static constexpr uint32_t VECTOR_SIZE = 128;
52 static constexpr uint32_t NUM4 = 4;45 static constexpr uint32_t NUM4 = 4;
53- 
54 CATLASS_DEVICE46 CATLASS_DEVICE
55 BlockEpilogue(Arch::Resource<ArchTag> &resource, uint32_t kvSplitCoreNum_ = 1)47 BlockEpilogue(Arch::Resource<ArchTag> &resource, uint32_t kvSplitCoreNum_ = 1)
56 {48 {
57 // Allocate UB space49 // Allocate UB space
58- constexpr uint32_t LO_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE_MLA;50+ constexpr uint32_t LO_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA;
59- constexpr uint32_t DM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 6 * UB_UINT8_LINE_SIZE;
60- constexpr uint32_t GL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 16 * UB_UINT8_LINE_SIZE;
61 constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE_MLA;51 constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE_MLA;
62- constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA;52+ constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA + 2 * UB_UINT8_LINE_SIZE;
63- constexpr uint32_t HM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE;53+ constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA + 7 * UB_UINT8_LINE_SIZE;
64- constexpr uint32_t GM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 14 * UB_UINT8_LINE_SIZE;54+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA + 9 * UB_UINT8_LINE_SIZE;
65- 55+ constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE;
56+ constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA + 5 * UB_UINT8_LINE_SIZE;
66 kvSplitCoreNum = kvSplitCoreNum_;57 kvSplitCoreNum = kvSplitCoreNum_;
67 loUbTensor = resource.ubBuf.template GetBufferByByte<float>(LO_UB_TENSOR_OFFSET);58 loUbTensor = resource.ubBuf.template GetBufferByByte<float>(LO_UB_TENSOR_OFFSET);
68 dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);59 dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
69- glUbTensor = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);60+ glUbTensor[0] = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
61+ glUbTensor[1] = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET + UB_UINT8_LINE_SIZE);
70 tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);62 tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
71 goUbTensor16 = resource.ubBuf.template GetBufferByByte<ElementOutput>(GO_UB_TENSOR_OFFSET);63 goUbTensor16 = resource.ubBuf.template GetBufferByByte<ElementOutput>(GO_UB_TENSOR_OFFSET);
72 goUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(GO_UB_TENSOR_OFFSET);64 goUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(GO_UB_TENSOR_OFFSET);
73 hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);65 hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);
74- gmUbTensor = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);66+ gmUbTensor[0] = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
67+ gmUbTensor[1] = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET + UB_UINT8_LINE_SIZE);
75 }68 }
76- 
77 CATLASS_DEVICE69 CATLASS_DEVICE
78 ~BlockEpilogue() {}70 ~BlockEpilogue() {}
79- 
80 CATLASS_DEVICE71 CATLASS_DEVICE
81 void SetMask(int32_t len)72 void SetMask(int32_t len)
82 {73 {
@@ -86,7 +77,6 @@ public:
86 for (int64_t i = 0; i < temp; i++) {77 for (int64_t i = 0; i < temp; i++) {
87 mask |= one << i;78 mask |= one << i;
88 }79 }
89- 
90 if (len == VECTOR_SIZE) {80 if (len == VECTOR_SIZE) {
91 AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);81 AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
92 } else if (len >= FLOAT_VECTOR_SIZE) {82 } else if (len >= FLOAT_VECTOR_SIZE) {
@@ -95,20 +85,18 @@ public:
95 AscendC::SetVectorMask<int8_t>(0x0, mask);85 AscendC::SetVectorMask<int8_t>(0x0, mask);
96 }86 }
97 }87 }
98- 
99 CATLASS_DEVICE88 CATLASS_DEVICE
100 void SetkvSplitCoreNum(uint32_t kvSplitCoreNum_)89 void SetkvSplitCoreNum(uint32_t kvSplitCoreNum_)
101 {90 {
102 kvSplitCoreNum = kvSplitCoreNum_;91 kvSplitCoreNum = kvSplitCoreNum_;
103 }92 }
104- 
105 CATLASS_DEVICE93 CATLASS_DEVICE
106 void SubCoreCompute(AscendC::GlobalTensor<ElementInput> gInput, AscendC::GlobalTensor<ElementUpdate> gUpdate,94 void SubCoreCompute(AscendC::GlobalTensor<ElementInput> gInput, AscendC::GlobalTensor<ElementUpdate> gUpdate,
107 AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementUpdate> gOCoreTmp,95 AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementUpdate> gOCoreTmp,
108 AscendC::GlobalTensor<ElementUpdate> gl, const LayoutInput &layoutInput,96 AscendC::GlobalTensor<ElementUpdate> gl, const LayoutInput &layoutInput,
109 const LayoutOutput &layoutOutput, const LayoutUpdate &layoutUpdate, uint32_t nIdx,97 const LayoutOutput &layoutOutput, const LayoutUpdate &layoutUpdate, uint32_t nIdx,
110 uint32_t isLastNTile, uint32_t needRowLoop, uint32_t rowLoopIdx, uint32_t rescaleOPingPongFlag,98 uint32_t isLastNTile, uint32_t needRowLoop, uint32_t rowLoopIdx, uint32_t rescaleOPingPongFlag,
111- uint32_t &glFlag)99+ uint32_t *glFlag, uint32_t taskPingPongFlag)
112 {100 {
113 uint32_t curRowNum = layoutInput.shape(0);101 uint32_t curRowNum = layoutInput.shape(0);
114 uint32_t embed = layoutInput.shape(1);102 uint32_t embed = layoutInput.shape(1);
@@ -173,11 +161,10 @@ public:
173 AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);161 AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
174 }162 }
175 AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(oPingPangFlag);163 AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(oPingPangFlag);
176- 
177 if (isLastNTile) {164 if (isLastNTile) {
178 // *** gl_block = expand_to_block(gl), 存放于 tv165 // *** gl_block = expand_to_block(gl), 存放于 tv
179 AscendC::Brcb(tvUbTensor.ReinterpretCast<uint32_t>(),166 AscendC::Brcb(tvUbTensor.ReinterpretCast<uint32_t>(),
180- glUbTensor.ReinterpretCast<uint32_t>()[rowLoopIdx * ROW_WISE_CYCLE_TILE],167+ glUbTensor[taskPingPongFlag].ReinterpretCast<uint32_t>()[rowLoopIdx * ROW_WISE_CYCLE_TILE],
181 curRowNumRound / FLOAT_BLOCK_SIZE,168 curRowNumRound / FLOAT_BLOCK_SIZE,
182 AscendC::BrcbRepeatParams(1, 8));169 AscendC::BrcbRepeatParams(1, 8));
183 AscendC::PipeBarrier<PIPE_V>();170 AscendC::PipeBarrier<PIPE_V>();
@@ -202,7 +189,6 @@ public:
202 AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);189 AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
203 }190 }
204 AscendC::PipeBarrier<PIPE_V>();191 AscendC::PipeBarrier<PIPE_V>();
205- 
206 if (kvSplitCoreNum != 1) {192 if (kvSplitCoreNum != 1) {
207 // log(l)193 // log(l)
208 AscendC::Ln<float, false>(194 AscendC::Ln<float, false>(
@@ -214,7 +200,7 @@ public:
214 AscendC::PipeBarrier<PIPE_V>();200 AscendC::PipeBarrier<PIPE_V>();
215 AscendC::Brcb(201 AscendC::Brcb(
216 hmUbTensor.ReinterpretCast<uint32_t>(),202 hmUbTensor.ReinterpretCast<uint32_t>(),
217- gmUbTensor.ReinterpretCast<uint32_t>()[rowLoopIdx * ROW_WISE_CYCLE_TILE],203+ gmUbTensor[taskPingPongFlag].ReinterpretCast<uint32_t>()[rowLoopIdx * ROW_WISE_CYCLE_TILE],
218 curRowNumRound / FLOAT_BLOCK_SIZE,204 curRowNumRound / FLOAT_BLOCK_SIZE,
219 AscendC::BrcbRepeatParams(1, 8));205 AscendC::BrcbRepeatParams(1, 8));
220 AscendC::PipeBarrier<PIPE_V>();206 AscendC::PipeBarrier<PIPE_V>();
@@ -227,21 +213,19 @@ public:
227 curRowNum,213 curRowNum,
228 AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));214 AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
229 AscendC::PipeBarrier<PIPE_V>();215 AscendC::PipeBarrier<PIPE_V>();
230- 
231 AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID2);216 AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID2);
232 AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID2);217 AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID2);
233 AscendC::DataCopyPad(gl, tvUbTensor,218 AscendC::DataCopyPad(gl, tvUbTensor,
234 AscendC::DataCopyExtParams(curRowNum, 4, 0, (kvSplitCoreNum - 1) * 4, 0));219 AscendC::DataCopyExtParams(curRowNum, 4, 0, (kvSplitCoreNum - 1) * 4, 0));
235- 220+ if (glFlag[taskPingPongFlag] == 0) {
236- if (glFlag == 0) {221+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(taskPingPongFlag + 2);
237- AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID2);222+ glFlag[taskPingPongFlag] = 1;
238- glFlag = 1;
239 }223 }
240 uint32_t srcGap = ((embed % 16 <= 8) && (embed % 16 > 0)) ? 1 : 0;224 uint32_t srcGap = ((embed % 16 <= 8) && (embed % 16 > 0)) ? 1 : 0;
241 AscendC::DataCopyPad(gOCoreTmp, goUbTensor32[oUbOffset],225 AscendC::DataCopyPad(gOCoreTmp, goUbTensor32[oUbOffset],
242 AscendC::DataCopyExtParams(curRowNum, embed * 4, srcGap, (kvSplitCoreNum - 1) * embed * 4, 0));226 AscendC::DataCopyExtParams(curRowNum, embed * 4, srcGap, (kvSplitCoreNum - 1) * embed * 4, 0));
243- AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID3);227+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
244- AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID3);228+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
245 AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID1);229 AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID1);
246 AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID1);230 AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID1);
247 } else {231 } else {
@@ -261,7 +245,6 @@ public:
261 }245 }
262 AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);246 AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
263 AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);247 AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
264- 
265 // ********************* move O to GM ************************248 // ********************* move O to GM ************************
266 AscendC::DataCopyPad(gOutput, goUbTensor16[oUbOffset * 2],249 AscendC::DataCopyPad(gOutput, goUbTensor16[oUbOffset * 2],
267 AscendC::DataCopyExtParams(curRowNum, embed * 2, 0, 0, 0));250 AscendC::DataCopyExtParams(curRowNum, embed * 2, 0, 0, 0));
@@ -275,25 +258,21 @@ public:
275 AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(oPingPangFlag + 4);258 AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(oPingPangFlag + 4);
276 oPingPangFlag = 1 - oPingPangFlag;259 oPingPangFlag = 1 - oPingPangFlag;
277 }260 }
278- 
279 CATLASS_DEVICE261 CATLASS_DEVICE
280 void operator()(AscendC::GlobalTensor<ElementInput> gInput, AscendC::GlobalTensor<ElementUpdate> gUpdate,262 void operator()(AscendC::GlobalTensor<ElementInput> gInput, AscendC::GlobalTensor<ElementUpdate> gUpdate,
281 AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementUpdate> gOCoreTmp,263 AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementUpdate> gOCoreTmp,
282 AscendC::GlobalTensor<ElementUpdate> gl, const LayoutInput &layoutInput,264 AscendC::GlobalTensor<ElementUpdate> gl, const LayoutInput &layoutInput,
283 const LayoutUpdate &layoutUpdate, const LayoutOutput &layoutOutput, GemmCoord actualBlockShape,265 const LayoutUpdate &layoutUpdate, const LayoutOutput &layoutOutput, GemmCoord actualBlockShape,
284- uint32_t nIdx, uint32_t isLastNTile, uint32_t rescaleOPingPongFlag, uint32_t &glFlag)266+ uint32_t nIdx, uint32_t isLastNTile, uint32_t rescaleOPingPongFlag, uint32_t *glFlag, uint32_t taskPingPongFlag)
285 {267 {
286 uint32_t embed = layoutInput.shape(1);268 uint32_t embed = layoutInput.shape(1);
287 uint32_t rowActual = actualBlockShape.m();269 uint32_t rowActual = actualBlockShape.m();
288 uint32_t columnActual = actualBlockShape.n();270 uint32_t columnActual = actualBlockShape.n();
289- 
290 uint32_t subBlockIdx = AscendC::GetSubBlockIdx();271 uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
291 uint32_t subBlockNum = AscendC::GetSubBlockNum();272 uint32_t subBlockNum = AscendC::GetSubBlockNum();
292- 
293 uint32_t curRowSplitSubBlock = rowActual / subBlockNum;273 uint32_t curRowSplitSubBlock = rowActual / subBlockNum;
294 uint32_t rowActualThisSubBlock = (subBlockIdx == 0) ? curRowSplitSubBlock : (rowActual - curRowSplitSubBlock);274 uint32_t rowActualThisSubBlock = (subBlockIdx == 0) ? curRowSplitSubBlock : (rowActual - curRowSplitSubBlock);
295 uint32_t rowOffsetSubBlock = subBlockIdx * curRowSplitSubBlock;275 uint32_t rowOffsetSubBlock = subBlockIdx * curRowSplitSubBlock;
296- 
297 if (rowActualThisSubBlock > 0) {276 if (rowActualThisSubBlock > 0) {
298 uint32_t rowLoop = (rowActualThisSubBlock + ROW_WISE_CYCLE_TILE - 1) / ROW_WISE_CYCLE_TILE;277 uint32_t rowLoop = (rowActualThisSubBlock + ROW_WISE_CYCLE_TILE - 1) / ROW_WISE_CYCLE_TILE;
299 uint32_t needRowLoop = (rowLoop > 1) ? 1 : 0;278 uint32_t needRowLoop = (rowLoop > 1) ? 1 : 0;
@@ -305,37 +284,31 @@ public:
305 int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetCurCycle, 0));284 int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetCurCycle, 0));
306 auto gInputThisCurCycle = gInput[offsetInput];285 auto gInputThisCurCycle = gInput[offsetInput];
307 auto layoutInputCurCycle = layoutInput.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual));286 auto layoutInputCurCycle = layoutInput.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual));
308- 
309 int64_t offsetUpdate = layoutUpdate.GetOffset(MatrixCoord(rowOffsetCurCycle, 0));287 int64_t offsetUpdate = layoutUpdate.GetOffset(MatrixCoord(rowOffsetCurCycle, 0));
310 auto gUpdateCurCycle = gUpdate[offsetUpdate];288 auto gUpdateCurCycle = gUpdate[offsetUpdate];
311 auto layoutUpdateCurCycle = layoutUpdate.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual));289 auto layoutUpdateCurCycle = layoutUpdate.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual));
312- 
313 int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetCurCycle, 0));290 int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetCurCycle, 0));
314 auto gOutputCurCycle = gOutput[offsetOutput];291 auto gOutputCurCycle = gOutput[offsetOutput];
315 auto layoutOutputCurCycle = layoutOutput.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual));292 auto layoutOutputCurCycle = layoutOutput.GetTileLayout(MatrixCoord(rowActualCurCycle, columnActual));
316- 
317 SubCoreCompute(gInputThisCurCycle, gUpdateCurCycle, gOutputCurCycle,293 SubCoreCompute(gInputThisCurCycle, gUpdateCurCycle, gOutputCurCycle,
318 gOCoreTmp[rowOffsetLoop * embed * kvSplitCoreNum],294 gOCoreTmp[rowOffsetLoop * embed * kvSplitCoreNum],
319 gl[rowOffsetLoop * kvSplitCoreNum], layoutInputCurCycle,295 gl[rowOffsetLoop * kvSplitCoreNum], layoutInputCurCycle,
320 layoutOutputCurCycle, layoutUpdateCurCycle, nIdx, isLastNTile, needRowLoop, rowLoopIdx,296 layoutOutputCurCycle, layoutUpdateCurCycle, nIdx, isLastNTile, needRowLoop, rowLoopIdx,
321- rescaleOPingPongFlag, glFlag);297+ rescaleOPingPongFlag, glFlag, taskPingPongFlag);
322 }298 }
323 }299 }
324 }300 }
325- 
326private:301private:
327 uint32_t kvSplitCoreNum = 1;302 uint32_t kvSplitCoreNum = 1;
328 uint32_t oPingPangFlag = 0;303 uint32_t oPingPangFlag = 0;
329 AscendC::LocalTensor<float> loUbTensor;304 AscendC::LocalTensor<float> loUbTensor;
330 AscendC::LocalTensor<float> dmUbTensor;305 AscendC::LocalTensor<float> dmUbTensor;
331- AscendC::LocalTensor<float> glUbTensor;306+ AscendC::LocalTensor<float> glUbTensor[2];
332 AscendC::LocalTensor<float> tvUbTensor;307 AscendC::LocalTensor<float> tvUbTensor;
333 AscendC::LocalTensor<ElementOutput> goUbTensor16;308 AscendC::LocalTensor<ElementOutput> goUbTensor16;
334 AscendC::LocalTensor<float> goUbTensor32;309 AscendC::LocalTensor<float> goUbTensor32;
335 AscendC::LocalTensor<float> hmUbTensor;310 AscendC::LocalTensor<float> hmUbTensor;
336- AscendC::LocalTensor<float> gmUbTensor;311+ AscendC::LocalTensor<float> gmUbTensor[2];
337};312};
338- 
339} // namespace Catlass::Epilogue::Block313} // namespace Catlass::Epilogue::Block
340- 
341#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_TP1_RESCALE_O_HPP314#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_MLA_TP1_RESCALE_O_HPP
Minclude/catlass/epilogue/block/block_epilogue_mla_tp1_softmax.hpp+54-34
@@ -55,13 +55,14 @@ public:
55 static constexpr uint32_t UB_UINT8_BLOCK_SIZE_MLA = 16384;55 static constexpr uint32_t UB_UINT8_BLOCK_SIZE_MLA = 16384;
56 static constexpr uint32_t VECTOR_SIZE = 128;56 static constexpr uint32_t VECTOR_SIZE = 128;
57 57 
58- static constexpr uint32_t REDUCE_UB_SIZE = 1024;58+ static constexpr uint32_t REDUCE_UB_SIZE = 1536;
59 static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32;59 static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32;
60 static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4;60 static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4;
61 static constexpr uint32_t S_BLOCK_STACK = 4;61 static constexpr uint32_t S_BLOCK_STACK = 4;
62 static constexpr int64_t UB_FLOAT_LINE_SIZE = 64;62 static constexpr int64_t UB_FLOAT_LINE_SIZE = 64;
63- static constexpr uint32_t M_SLICE = 16;63+ static constexpr uint32_t M_SLICE = 24;
64 static constexpr uint32_t QK_READY_ID = 1;64 static constexpr uint32_t QK_READY_ID = 1;
65+ static constexpr uint32_t LS_CHUNK_SIZE = 12288;
65 66 
66 CATLASS_DEVICE67 CATLASS_DEVICE
67 BlockEpilogue(Arch::Resource<ArchTag> &resource, half tor_, uint32_t kvSplitCoreNum_ = 1)68 BlockEpilogue(Arch::Resource<ArchTag> &resource, half tor_, uint32_t kvSplitCoreNum_ = 1)
@@ -69,13 +70,13 @@ public:
69 // Allocate UB space70 // Allocate UB space
70 constexpr uint32_t LS_UB_TENSOR_OFFSET = 0;71 constexpr uint32_t LS_UB_TENSOR_OFFSET = 0;
71 constexpr uint32_t LP_UB_TENSOR_OFFSET = 0;72 constexpr uint32_t LP_UB_TENSOR_OFFSET = 0;
72- constexpr uint32_t LM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA;73+ constexpr uint32_t LM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA;
73- constexpr uint32_t HM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE;74+ constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA + 1 * UB_UINT8_LINE_SIZE;
74- constexpr uint32_t DM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 6 * UB_UINT8_LINE_SIZE;75+ constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA + 2 * UB_UINT8_LINE_SIZE;
75- constexpr uint32_t LL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 10 * UB_UINT8_LINE_SIZE;76+ constexpr uint32_t LL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA + 4 * UB_UINT8_LINE_SIZE;
76- constexpr uint32_t GM_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 14 * UB_UINT8_LINE_SIZE;77+ constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA + 5 * UB_UINT8_LINE_SIZE;
77- constexpr uint32_t GL_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE_MLA + 16 * UB_UINT8_LINE_SIZE;78+ constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA + 7 * UB_UINT8_LINE_SIZE;
78- constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA;79+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE_MLA + 9 * UB_UINT8_LINE_SIZE;
79 80 
80 tor = tor_;81 tor = tor_;
81 kvSplitCoreNum = kvSplitCoreNum_;82 kvSplitCoreNum = kvSplitCoreNum_;
@@ -83,11 +84,13 @@ public:
83 lpUbTensor = resource.ubBuf.template GetBufferByByte<ElementOutput>(LP_UB_TENSOR_OFFSET);84 lpUbTensor = resource.ubBuf.template GetBufferByByte<ElementOutput>(LP_UB_TENSOR_OFFSET);
84 lmUbTensor = resource.ubBuf.template GetBufferByByte<float>(LM_UB_TENSOR_OFFSET);85 lmUbTensor = resource.ubBuf.template GetBufferByByte<float>(LM_UB_TENSOR_OFFSET);
85 hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);86 hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);
86- gmUbTensor = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);87+ gmUbTensor[0] = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
88+ gmUbTensor[1] = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET + UB_UINT8_LINE_SIZE);
87 dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);89 dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
88 llUbTensor = resource.ubBuf.template GetBufferByByte<float>(LL_UB_TENSOR_OFFSET);90 llUbTensor = resource.ubBuf.template GetBufferByByte<float>(LL_UB_TENSOR_OFFSET);
89 tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);91 tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
90- glUbTensor = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);92+ glUbTensor[0] = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
93+ glUbTensor[1] = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET + UB_UINT8_LINE_SIZE);
91 }94 }
92 95 
93 CATLASS_DEVICE96 CATLASS_DEVICE
@@ -315,13 +318,14 @@ public:
315 CATLASS_DEVICE318 CATLASS_DEVICE
316 void SubCoreCompute(AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,319 void SubCoreCompute(AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
317 uint32_t m, uint32_t nReal, uint32_t nStride, uint32_t pingpongFlag, uint32_t rowOffset,320 uint32_t m, uint32_t nReal, uint32_t nStride, uint32_t pingpongFlag, uint32_t rowOffset,
318- uint32_t sUbOffset, uint32_t nIdx, uint32_t &glFlag)321+ uint32_t sUbOffset, uint32_t nIdx, uint32_t *glFlag, uint32_t taskPingPongFlag, uint32_t gSPingPongFlag)
319 {322 {
320 uint32_t round_m = (m + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE * FLOAT_BLOCK_SIZE;323 uint32_t round_m = (m + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE * FLOAT_BLOCK_SIZE;
321- AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(pingpongFlag);324+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(gSPingPongFlag);
322 // input QK325 // input QK
323 AscendC::DataCopy(lsUbTensor[sUbOffset], gInput, AscendC::DataCopyParams(m, nStride / FLOAT_BLOCK_SIZE, 0, 0));326 AscendC::DataCopy(lsUbTensor[sUbOffset], gInput, AscendC::DataCopyParams(m, nStride / FLOAT_BLOCK_SIZE, 0, 0));
324 327 
328+ 
325 AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);329 AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
326 AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);330 AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
327 331 
@@ -334,9 +338,9 @@ public:
334 338 
335 if (kvSplitCoreNum != 1) {339 if (kvSplitCoreNum != 1) {
336 if (nIdx == 0) {340 if (nIdx == 0) {
337- if (glFlag == 1) {341+ if (glFlag[taskPingPongFlag] == 1) {
338- AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID2);342+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(taskPingPongFlag + 2);
339- glFlag = 0;343+ glFlag[taskPingPongFlag] = 0;
340 }344 }
341 }345 }
342 }346 }
@@ -349,45 +353,51 @@ public:
349 RowmaxTAILTILE(lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride);353 RowmaxTAILTILE(lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride);
350 }354 }
351 355 
356+ 
352 if (nIdx == 0) {357 if (nIdx == 0) {
353 AscendC::DataCopy(hmUbTensor[rowOffset], lmUbTensor[rowOffset],358 AscendC::DataCopy(hmUbTensor[rowOffset], lmUbTensor[rowOffset],
354 AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0));359 AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0));
360+ 
355 AscendC::PipeBarrier<PIPE_V>();361 AscendC::PipeBarrier<PIPE_V>();
356 } else {362 } else {
357 SetVecMask(m);363 SetVecMask(m);
358 // *** hm = vmax(lm, gm)364 // *** hm = vmax(lm, gm)
359- AscendC::Max<float, false>(hmUbTensor[rowOffset], lmUbTensor[rowOffset], gmUbTensor[rowOffset], (uint64_t)0,365+ AscendC::Max<float, false>(hmUbTensor[rowOffset], lmUbTensor[rowOffset], gmUbTensor[taskPingPongFlag][rowOffset], (uint64_t)0,
360 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));366 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
361 367 
362 AscendC::PipeBarrier<PIPE_V>();368 AscendC::PipeBarrier<PIPE_V>();
363 // *** dm = gm - hm369 // *** dm = gm - hm
364- AscendC::Sub<float, false>(dmUbTensor[((nIdx / S_BLOCK_STACK) % 2) * UB_FLOAT_LINE_SIZE + rowOffset],370+ AscendC::Sub<float, false>(dmUbTensor[gSPingPongFlag * UB_FLOAT_LINE_SIZE + rowOffset],
365- gmUbTensor[rowOffset], hmUbTensor[rowOffset], (uint64_t)0, 1,371+ gmUbTensor[taskPingPongFlag][rowOffset], hmUbTensor[rowOffset], (uint64_t)0, 1,
366 AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));372 AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
367 373 
368 AscendC::PipeBarrier<PIPE_V>();374 AscendC::PipeBarrier<PIPE_V>();
369 // *** dm = exp(dm)375 // *** dm = exp(dm)
370- AscendC::Exp<float, false>(dmUbTensor[((nIdx / S_BLOCK_STACK) % 2) * UB_FLOAT_LINE_SIZE + rowOffset],376+ AscendC::Exp<float, false>(dmUbTensor[gSPingPongFlag * UB_FLOAT_LINE_SIZE + rowOffset],
371- dmUbTensor[((nIdx / S_BLOCK_STACK) % 2) * UB_FLOAT_LINE_SIZE + rowOffset],377+ dmUbTensor[gSPingPongFlag * UB_FLOAT_LINE_SIZE + rowOffset],
372 (uint64_t)0, 1, AscendC::UnaryRepeatParams(1, 1, 8, 8));378 (uint64_t)0, 1, AscendC::UnaryRepeatParams(1, 1, 8, 8));
373 }379 }
380+ 
374 AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);381 AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
375 AscendC::PipeBarrier<PIPE_V>();382 AscendC::PipeBarrier<PIPE_V>();
376 // *** gm = hm383 // *** gm = hm
377- AscendC::DataCopy(gmUbTensor[rowOffset], hmUbTensor[rowOffset],384+ AscendC::DataCopy(gmUbTensor[taskPingPongFlag][rowOffset], hmUbTensor[rowOffset],
378 AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0));385 AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0));
379 AscendC::PipeBarrier<PIPE_V>();386 AscendC::PipeBarrier<PIPE_V>();
387+ 
380 // *** hm_block = expand_to_block(hm), 存放于 tv388 // *** hm_block = expand_to_block(hm), 存放于 tv
381 AscendC::Brcb(tvUbTensor.template ReinterpretCast<uint32_t>(),389 AscendC::Brcb(tvUbTensor.template ReinterpretCast<uint32_t>(),
382 hmUbTensor[rowOffset].template ReinterpretCast<uint32_t>(), round_m / FLOAT_BLOCK_SIZE,390 hmUbTensor[rowOffset].template ReinterpretCast<uint32_t>(), round_m / FLOAT_BLOCK_SIZE,
383 AscendC::BrcbRepeatParams(1, 8));391 AscendC::BrcbRepeatParams(1, 8));
384 AscendC::PipeBarrier<PIPE_V>();392 AscendC::PipeBarrier<PIPE_V>();
393+ 
385 // *** ls = ls - hm_block394 // *** ls = ls - hm_block
386 for (uint32_t subIdx = 0; subIdx < nReal / FLOAT_VECTOR_SIZE; ++subIdx) {395 for (uint32_t subIdx = 0; subIdx < nReal / FLOAT_VECTOR_SIZE; ++subIdx) {
387 AscendC::Sub<float, false>(396 AscendC::Sub<float, false>(
388 lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE],397 lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE],
389 tvUbTensor, (uint64_t)0, m,398 tvUbTensor, (uint64_t)0, m,
390 AscendC::BinaryRepeatParams(1, 1, 0, nStride / FLOAT_BLOCK_SIZE, nStride / FLOAT_BLOCK_SIZE, 1));399 AscendC::BinaryRepeatParams(1, 1, 0, nStride / FLOAT_BLOCK_SIZE, nStride / FLOAT_BLOCK_SIZE, 1));
400+ 
391 }401 }
392 if (nReal % FLOAT_VECTOR_SIZE > 0) {402 if (nReal % FLOAT_VECTOR_SIZE > 0) {
393 SetVecMask(nReal % FLOAT_VECTOR_SIZE);403 SetVecMask(nReal % FLOAT_VECTOR_SIZE);
@@ -396,6 +406,7 @@ public:
396 lsUbTensor[sUbOffset][nReal / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0, m,406 lsUbTensor[sUbOffset][nReal / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0, m,
397 AscendC::BinaryRepeatParams(1, 1, 0, nStride / FLOAT_BLOCK_SIZE, nStride / FLOAT_BLOCK_SIZE, 1));407 AscendC::BinaryRepeatParams(1, 1, 0, nStride / FLOAT_BLOCK_SIZE, nStride / FLOAT_BLOCK_SIZE, 1));
398 AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);408 AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
409+ 
399 }410 }
400 AscendC::PipeBarrier<PIPE_V>();411 AscendC::PipeBarrier<PIPE_V>();
401 412 
@@ -404,6 +415,7 @@ public:
404 (m * nStride + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,415 (m * nStride + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
405 AscendC::UnaryRepeatParams(1, 1, 8, 8));416 AscendC::UnaryRepeatParams(1, 1, 8, 8));
406 AscendC::PipeBarrier<PIPE_V>();417 AscendC::PipeBarrier<PIPE_V>();
418+ 
407 // *** ll = rowsum(ls32)419 // *** ll = rowsum(ls32)
408 if (nReal == 512) {420 if (nReal == 512) {
409 RowsumSPECTILE512(lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride);421 RowsumSPECTILE512(lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride);
@@ -413,6 +425,7 @@ public:
413 RowsumTAILTILE(lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride);425 RowsumTAILTILE(lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, round_m, nReal, nStride);
414 }426 }
415 427 
428+ 
416 // *** lp = castfp32to16(ls)429 // *** lp = castfp32to16(ls)
417 if (std::is_same<ElementOutput, bfloat16_t>::value) {430 if (std::is_same<ElementOutput, bfloat16_t>::value) {
418 AscendC::Cast<ElementOutput, float, false>(431 AscendC::Cast<ElementOutput, float, false>(
@@ -426,24 +439,28 @@ public:
426 439 
427 AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);440 AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
428 AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);441 AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
442+ 
429 AscendC::DataCopy(gOutput, lpUbTensor[sUbOffset * 2], AscendC::DataCopyParams(m, nStride * 2 / 32, 0, 0));443 AscendC::DataCopy(gOutput, lpUbTensor[sUbOffset * 2], AscendC::DataCopyParams(m, nStride * 2 / 32, 0, 0));
430- AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(pingpongFlag);444+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(gSPingPongFlag);
431 if (nIdx == 0) {445 if (nIdx == 0) {
432 // *** gl = ll446 // *** gl = ll
433- AscendC::DataCopy(glUbTensor[rowOffset], llUbTensor[rowOffset],447+ AscendC::DataCopy(glUbTensor[taskPingPongFlag][rowOffset], llUbTensor[rowOffset],
434 AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0));448 AscendC::DataCopyParams(1, round_m / FLOAT_BLOCK_SIZE, 0, 0));
449+ 
435 AscendC::PipeBarrier<PIPE_V>();450 AscendC::PipeBarrier<PIPE_V>();
436 } else {451 } else {
437 SetVecMask(m);452 SetVecMask(m);
438- // *** gl = dm * gl453+ // // *** gl = dm * gl
439 AscendC::Mul<float, false>(454 AscendC::Mul<float, false>(
440- glUbTensor[rowOffset], dmUbTensor[((nIdx / S_BLOCK_STACK) % 2) * UB_FLOAT_LINE_SIZE + rowOffset],455+ glUbTensor[taskPingPongFlag][rowOffset], dmUbTensor[gSPingPongFlag * UB_FLOAT_LINE_SIZE + rowOffset],
441- glUbTensor[rowOffset], (uint64_t)0, 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));456+ glUbTensor[taskPingPongFlag][rowOffset], (uint64_t)0, 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
442 AscendC::PipeBarrier<PIPE_V>();457 AscendC::PipeBarrier<PIPE_V>();
443- // *** gl = ll + gl458+ 
444- AscendC::Add<float, false>(glUbTensor[rowOffset], glUbTensor[rowOffset], llUbTensor[rowOffset], (uint64_t)0,459+ // // *** gl = ll + gl
460+ AscendC::Add<float, false>(glUbTensor[taskPingPongFlag][rowOffset], glUbTensor[taskPingPongFlag][rowOffset], llUbTensor[rowOffset], (uint64_t)0,
445 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));461 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
446 AscendC::PipeBarrier<PIPE_V>();462 AscendC::PipeBarrier<PIPE_V>();
463+ 
447 AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);464 AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
448 }465 }
449 }466 }
@@ -451,11 +468,12 @@ public:
451 CATLASS_DEVICE468 CATLASS_DEVICE
452 void operator()(AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,469 void operator()(AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
453 const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, GemmCoord actualBlockShape,470 const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, GemmCoord actualBlockShape,
454- uint32_t nIdx, uint32_t &glFlag)471+ uint32_t nIdx, uint32_t *glFlag, uint32_t taskPingPongFlag, uint32_t gSPingPongFlag)
455 {472 {
456 uint32_t cur_head_num = actualBlockShape.m();473 uint32_t cur_head_num = actualBlockShape.m();
457 uint32_t qkN = actualBlockShape.n();474 uint32_t qkN = actualBlockShape.n();
458 uint32_t qkRoundN = layoutInput.stride(0);475 uint32_t qkRoundN = layoutInput.stride(0);
476+ 
459 uint32_t pingpongFlag = 0;477 uint32_t pingpongFlag = 0;
460 478 
461 uint32_t subBlockIdx = AscendC::GetSubBlockIdx();479 uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
@@ -468,7 +486,7 @@ public:
468 for (uint32_t mInd = 0; mInd < mEnd; mInd++) {486 for (uint32_t mInd = 0; mInd < mEnd; mInd++) {
469 uint32_t rowOffset = mInd * M_SLICE;487 uint32_t rowOffset = mInd * M_SLICE;
470 uint32_t currM = mInd == mEnd - 1 ? subM - rowOffset : M_SLICE;488 uint32_t currM = mInd == mEnd - 1 ? subM - rowOffset : M_SLICE;
471- uint32_t sUbOffset = pingpongFlag * 8192;489+ uint32_t sUbOffset = gSPingPongFlag * LS_CHUNK_SIZE;
472 int64_t offsetOutput = rowOffset * qkRoundN;490 int64_t offsetOutput = rowOffset * qkRoundN;
473 auto gOutputThisSubBlock = gOutput[offsetOutput];491 auto gOutputThisSubBlock = gOutput[offsetOutput];
474 int64_t offsetInput = rowOffset * qkRoundN;492 int64_t offsetInput = rowOffset * qkRoundN;
@@ -480,8 +498,9 @@ public:
480 continue;498 continue;
481 }499 }
482 SubCoreCompute(gOutputThisSubBlock, gInputThisSubBlock, currM, qkN, qkRoundN, pingpongFlag, rowOffset,500 SubCoreCompute(gOutputThisSubBlock, gInputThisSubBlock, currM, qkN, qkRoundN, pingpongFlag, rowOffset,
483- sUbOffset, nIdx, glFlag);501+ sUbOffset, nIdx, glFlag, taskPingPongFlag, gSPingPongFlag);
484 pingpongFlag = 1 - pingpongFlag;502 pingpongFlag = 1 - pingpongFlag;
503+ 
485 }504 }
486 }505 }
487 506 
@@ -494,13 +513,14 @@ private:
494 AscendC::LocalTensor<ElementOutput> lpUbTensor;513 AscendC::LocalTensor<ElementOutput> lpUbTensor;
495 AscendC::LocalTensor<float> lmUbTensor;514 AscendC::LocalTensor<float> lmUbTensor;
496 AscendC::LocalTensor<float> hmUbTensor;515 AscendC::LocalTensor<float> hmUbTensor;
497- AscendC::LocalTensor<float> gmUbTensor;516+ AscendC::LocalTensor<float> gmUbTensor[2];
498 AscendC::LocalTensor<float> dmUbTensor;517 AscendC::LocalTensor<float> dmUbTensor;
499 AscendC::LocalTensor<float> llUbTensor;518 AscendC::LocalTensor<float> llUbTensor;
500 AscendC::LocalTensor<float> tvUbTensor;519 AscendC::LocalTensor<float> tvUbTensor;
501- AscendC::LocalTensor<float> glUbTensor;520+ AscendC::LocalTensor<float> glUbTensor[2];
502 521 
503 Arch::CrossCoreFlag qkReady{QK_READY_ID};522 Arch::CrossCoreFlag qkReady{QK_READY_ID};
523+ 
504};524};
505 525 
506} // namespace Catlass::Epilogue::Block526} // namespace Catlass::Epilogue::Block
Minclude/catlass/epilogue/dispatch_policy.hpp+10-0
@@ -91,6 +91,16 @@ struct EpilogueAtlasA2MLATP1RescaleO {
91 using ArchTag = Arch::AtlasA2;91 using ArchTag = Arch::AtlasA2;
92};92};
93 93 
94+// For AtlasA2, AMLA TP1 Softmax
95+struct EpilogueAtlasA2AMLATP1Softmax {
96+ using ArchTag = Arch::AtlasA2;
97+};
98+ 
99+// For AtlasA2, AMLA TP1 RescaleO
100+struct EpilogueAtlasA2AMLATP1RescaleO {
101+ using ArchTag = Arch::AtlasA2;
102+};
103+ 
94// For AtlasA2, per token dequant104// For AtlasA2, per token dequant
95template <uint32_t UB_STAGES_>105template <uint32_t UB_STAGES_>
96struct EpilogueAtlasA2PerTokenDequant {106struct EpilogueAtlasA2PerTokenDequant {
Minclude/catlass/gemm/block/block_mmad.hpp+4-1
@@ -33,7 +33,7 @@ struct BlockMmad {
33 static_assert(DEPENDENT_FALSE<DispatchPolicy>, "BlockMmad is not implemented for this DispatchPolicy");33 static_assert(DEPENDENT_FALSE<DispatchPolicy>, "BlockMmad is not implemented for this DispatchPolicy");
34};34};
35 35 
36-#if (defined (CATLASS_ARCH) && CATLASS_ARCH == 2201) 36+#if (defined (CATLASS_ARCH) && CATLASS_ARCH == 2201)
37/// new add for the reason that i am using the dispatchpolicy which is same as the policy of the optimized_matmul37/// new add for the reason that i am using the dispatchpolicy which is same as the policy of the optimized_matmul
38// so i add a new one class to avoid the conflict38// so i add a new one class to avoid the conflict
39template <39template <
@@ -137,6 +137,9 @@ struct BlockPrologue {
137#include "catlass/gemm/block/block_mmad_mla_pv.hpp"137#include "catlass/gemm/block/block_mmad_mla_pv.hpp"
138#include "catlass/gemm/block/block_mmad_mla_qk_tp1_spec.hpp"138#include "catlass/gemm/block/block_mmad_mla_qk_tp1_spec.hpp"
139#include "catlass/gemm/block/block_mmad_mla_pv_tp1_spec.hpp"139#include "catlass/gemm/block/block_mmad_mla_pv_tp1_spec.hpp"
140+ 
141+#include "catlass/gemm/block/block_mmad_amla_pv_tp1_spec.hpp"
142+ 
140#include "catlass/gemm/block/block_mmad_preload.hpp"143#include "catlass/gemm/block/block_mmad_preload.hpp"
141#include "catlass/gemm/block/block_mmad_preload_async.hpp"144#include "catlass/gemm/block/block_mmad_preload_async.hpp"
142#include "catlass/gemm/block/block_mmad_preload_async_with_callback.hpp"145#include "catlass/gemm/block/block_mmad_preload_async_with_callback.hpp"