Star86
139
代码
代码
Issues122
Pull Requests31
流水线
项目讨论
Wiki
分析
Star86
139
已合并
[Feature][ops]迁移fia里pfa的op_host到MindIE-SD #410
yyansifu创建于 7月3日
[Feature][ops]迁移fia里pfa的op_host到MindIE-SD #410
已合并
yyansifu创建于 7月3日
从yyansifu/MindIE-SD_pfadev合入到Ascend/MindIE-SDdev
共 8 个文件变更+14189-0
Acsrc/ops/prompt_flash_attention/op_host/prompt_flash_attention_tiling.cpp+7070-0
Acsrc/ops/prompt_flash_attention/op_host/prompt_flash_attention_tiling.h+880-0
@@ -0,0 +1,880 @@
1+/**
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
3+ * MindIE is licensed under Mulan PSL v2.
4+ * You can use this software according to the terms and conditions of the Mulan PSL v2.
5+ * You may obtain a copy of Mulan PSL v2 at:
6+ * http://license.coscl.org.cn/MulanPSL2
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
8+ * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
9+ * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
10+ * See the Mulan PSL v2 for more details.
11+ */
12+ 
13+/*!
14+ * \file prompt_flash_attention_tiling.h
15+ * \brief
16+ */
17+#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_PROMPTFLASHATTENTION_H_
18+#define AIR_CXX_RUNTIME_V2_OP_IMPL_PROMPTFLASHATTENTION_H_
19+#include "exe_graph/runtime/tiling_context.h"
20+#include "op_host/data_copy_transpose_tiling_def.h"
21+#include "op_host/data_copy_transpose_tiling.h"
22+#include "register/tilingdata_base.h"
23+#include "tiling/tiling_api.h"
24+#include "register/op_def_registry.h"
25+#ifdef ASCENDC_OP_TEST
26+#define PFA_EXTERN_C extern "C"
27+#else
28+#define PFA_EXTERN_C
29+#endif
30+ 
31+#include "prompt_flash_attention_tiling_compile_info.h"
32+#include "prompt_flash_attention_tiling_const.h"
33+#include "prompt_flash_attention_tiling_context.h"
34+#include "prompt_flash_attention_tiling_struct.h"
35+#include "../../common/op_host/fia_tiling_templates_registry.h"
36+ 
37+namespace optiling {
38+ 
39+struct PFAShapeInfo {
40+ uint32_t b = 0;
41+ uint64_t n = 0;
42+ uint64_t s = 0;
43+ uint32_t d = 0;
44+ uint64_t h = 0;
45+ uint32_t t = 0;
46+};
47+ 
48+BEGIN_TILING_DATA_DEF(PromptAttentionBaseParams)
49+TILING_DATA_FIELD_DEF(uint32_t, batchSize);
50+TILING_DATA_FIELD_DEF(uint64_t, headNumSize);
51+TILING_DATA_FIELD_DEF(uint32_t, seqSize);
52+TILING_DATA_FIELD_DEF(uint64_t, headSize);
53+TILING_DATA_FIELD_DEF(float, scaleValue);
54+TILING_DATA_FIELD_DEF(int32_t, preTokens);
55+TILING_DATA_FIELD_DEF(int32_t, nextTokens);
56+TILING_DATA_FIELD_DEF(int32_t, blockSize);
57+TILING_DATA_FIELD_DEF(int32_t, blockTableDim2);
58+TILING_DATA_FIELD_DEF(int32_t, PABlockNumSum);
59+TILING_DATA_FIELD_DEF(uint32_t, dimNumOfseq);
60+TILING_DATA_FIELD_DEF(uint32_t, typeByteNum);
61+TILING_DATA_FIELD_DEF(uint32_t, seqInnerSize);
62+TILING_DATA_FIELD_DEF(uint32_t, prefixSeqInnerSize);
63+TILING_DATA_FIELD_DEF(uint32_t, usePseShift);
64+TILING_DATA_FIELD_DEF(uint32_t, useMask);
65+TILING_DATA_FIELD_DEF(uint64_t, headNumRatio);
66+TILING_DATA_FIELD_DEF(uint32_t, attenMaskElemType);
67+TILING_DATA_FIELD_DEF(uint32_t, pseShiftTypeByteNum);
68+TILING_DATA_FIELD_DEF(uint32_t, pseMaskMaxSize);
69+TILING_DATA_FIELD_DEF(uint32_t, maskTypeByteNum);
70+TILING_DATA_FIELD_DEF(uint32_t, outputTypeByteNum);
71+TILING_DATA_FIELD_DEF(uint32_t, softmaxTypeByteNum);
72+TILING_DATA_FIELD_DEF(uint32_t, sparseMode);
73+TILING_DATA_FIELD_DEF(uint32_t, alignedHeadSize);
74+TILING_DATA_FIELD_DEF(uint32_t, splitS2);
75+TILING_DATA_FIELD_DEF(uint32_t, splitD);
76+TILING_DATA_FIELD_DEF(uint32_t, layoutType);
77+TILING_DATA_FIELD_DEF(uint32_t, PAlayoutType);
78+TILING_DATA_FIELD_DEF(uint32_t, pseShiftS1Size);
79+TILING_DATA_FIELD_DEF(uint32_t, pseShiftS2Size);
80+TILING_DATA_FIELD_DEF(uint32_t, maskKVsSize);
81+TILING_DATA_FIELD_DEF(uint32_t, maskQsSize);
82+TILING_DATA_FIELD_DEF(uint32_t, isLayoutSH);
83+TILING_DATA_FIELD_DEF(uint32_t, isActualSeqLengthsNull);
84+TILING_DATA_FIELD_DEF(uint32_t, isActualSeqLengthsKVNull);
85+TILING_DATA_FIELD_DEF(uint32_t, actualSeqLengthsSize);
86+TILING_DATA_FIELD_DEF(uint32_t, actualSeqLengthsKVSize);
87+TILING_DATA_FIELD_DEF(uint32_t, deqScaleFlag);
88+TILING_DATA_FIELD_DEF(uint32_t, deqScale2Flag);
89+TILING_DATA_FIELD_DEF(uint32_t, isAntiPerchannel);
90+TILING_DATA_FIELD_DEF(uint32_t, isRowInvalid);
91+TILING_DATA_FIELD_DEF(uint32_t, softmaxOuterSize);
92+TILING_DATA_FIELD_DEF(uint32_t, isQuant2Perchannel);
93+TILING_DATA_FIELD_DEF(uint32_t, isQuant2BF16);
94+TILING_DATA_FIELD_DEF(uint32_t, isKvContinuous);
95+TILING_DATA_FIELD_DEF(uint32_t, fromFused);
96+TILING_DATA_FIELD_DEF(uint32_t, isBSNDOut);
97+TILING_DATA_FIELD_DEF(uint32_t, isIFA);
98+TILING_DATA_FIELD_DEF(uint32_t, isSoftMaxLseEnable);
99+TILING_DATA_FIELD_DEF(uint32_t, isActualSharedPrefixLenNull);
100+TILING_DATA_FIELD_DEF(uint32_t, isQHasLeftPadding);
101+TILING_DATA_FIELD_DEF(uint32_t, isKVHasLeftPadding);
102+TILING_DATA_FIELD_DEF(int64_t, keyAntiquantMode);
103+TILING_DATA_FIELD_DEF(int64_t, valueAntiquantMode);
104+TILING_DATA_FIELD_DEF(uint32_t, hasKeyAntiquantOffset);
105+TILING_DATA_FIELD_DEF(uint32_t, isMsd);
106+TILING_DATA_FIELD_DEF(uint32_t, isQuant2FP16);
107+TILING_DATA_FIELD_DEF(uint32_t, ropeHeadSize);
108+TILING_DATA_FIELD_DEF(uint32_t, qkHeadSize);
109+TILING_DATA_FIELD_DEF(uint32_t, vHeadSize);
110+TILING_DATA_FIELD_DEF(uint32_t, gOfMla);
111+END_TILING_DATA_DEF;
112+REGISTER_TILING_DATA_CLASS(PromptAttentionBaseParamsOp, PromptAttentionBaseParams)
113+ 
114+BEGIN_TILING_DATA_DEF(PromptAttentionBaseApiBaseParams)
115+TILING_DATA_FIELD_DEF(uint32_t, batchSize);
116+TILING_DATA_FIELD_DEF(uint64_t, headNumSize);
117+TILING_DATA_FIELD_DEF(uint64_t, headSize);
118+TILING_DATA_FIELD_DEF(uint32_t, maskTypeByteNum);
119+ 
120+TILING_DATA_FIELD_DEF(uint32_t, inputLayoutType);
121+TILING_DATA_FIELD_DEF(uint64_t, kvHeadNumSize);
122+TILING_DATA_FIELD_DEF(uint32_t, maxSeqLen);
123+TILING_DATA_FIELD_DEF(uint32_t, maxKvSeqLen);
124+TILING_DATA_FIELD_DEF(uint32_t, totalQBlkNum);
125+TILING_DATA_FIELD_DEF(uint32_t, embeddingSizeV);
126+TILING_DATA_FIELD_DEF(uint32_t, quantType);
127+TILING_DATA_FIELD_DEF(uint32_t, dataShapeType);
128+TILING_DATA_FIELD_DEF(uint32_t, scaleType);
129+TILING_DATA_FIELD_DEF(uint64_t, workSize);
130+TILING_DATA_FIELD_DEF(float, tor);
131+TILING_DATA_FIELD_DEF(uint32_t, headStride);
132+TILING_DATA_FIELD_DEF(uint32_t, maskStride);
133+TILING_DATA_FIELD_DEF(uint32_t, isTriuMask);
134+TILING_DATA_FIELD_DEF(uint32_t, isClamp);
135+TILING_DATA_FIELD_DEF(uint32_t, clampMin);
136+TILING_DATA_FIELD_DEF(uint32_t, clampMax);
137+TILING_DATA_FIELD_DEF(uint32_t, tilingHeadSize);
138+TILING_DATA_FIELD_DEF(uint32_t, tilingParaSize);
139+TILING_DATA_FIELD_DEF(uint32_t, isLongSeq);
140+TILING_DATA_FIELD_DEF(uint32_t, isAlibiMaskSqrt);
141+TILING_DATA_FIELD_DEF(uint32_t, maskType);
142+TILING_DATA_FIELD_DEF(uint32_t, alibiCompressOffset);
143+TILING_DATA_FIELD_DEF(uint32_t, alibiLeftAlign);
144+TILING_DATA_FIELD_DEF(uint32_t, ppMScalar);
145+TILING_DATA_FIELD_DEF(uint32_t, ppNScalar);
146+TILING_DATA_FIELD_DEF(uint32_t, totalQBlkNumFirst);
147+END_TILING_DATA_DEF;
148+REGISTER_TILING_DATA_CLASS(PromptAttentionBaseApiBaseParamsOp, PromptAttentionBaseApiBaseParams)
149+ 
150+BEGIN_TILING_DATA_DEF(PromptAttentionSeqParams)
151+// Temporary reuse
152+TILING_DATA_FIELD_DEF_ARR(uint32_t, 64, CoreHeadNumTail); // coreNStart
153+TILING_DATA_FIELD_DEF_ARR(uint32_t, 64, actualS1); // coreNEnd
154+TILING_DATA_FIELD_DEF_ARR(uint32_t, 64, actualCoreNums); // coreSidStart
155+TILING_DATA_FIELD_DEF_ARR(uint32_t, 64, singleCoreHeadNumSize); // coreSidEnd
156+TILING_DATA_FIELD_DEF_ARR(uint32_t, 64, coreSeqPosStart);
157+TILING_DATA_FIELD_DEF_ARR(uint32_t, 64, coreSeqPosEnd);
158+END_TILING_DATA_DEF;
159+REGISTER_TILING_DATA_CLASS(PromptAttentionSeqParamsOp, PromptAttentionSeqParams)
160+ 
161+BEGIN_TILING_DATA_DEF(PromptAttentionSplitCoreParams)
162+TILING_DATA_FIELD_DEF_ARR(uint32_t, 50, startBlkArray);
163+TILING_DATA_FIELD_DEF_ARR(uint32_t, 50, endBlkArray);
164+END_TILING_DATA_DEF;
165+REGISTER_TILING_DATA_CLASS(PromptAttentionSplitCoreParamsOp, PromptAttentionSplitCoreParams);
166+ 
167+BEGIN_TILING_DATA_DEF(PromptAttentionSingleCoreParams)
168+TILING_DATA_FIELD_DEF(uint32_t, singleProcessSInnerSize);
169+TILING_DATA_FIELD_DEF(uint32_t, singleProcessSOuterSize);
170+TILING_DATA_FIELD_DEF(uint32_t, multiSmaxsInnerLoopTimes);
171+TILING_DATA_FIELD_DEF(uint32_t, actualCoreNums);
172+TILING_DATA_FIELD_DEF(uint32_t, pseShiftBatch);
173+TILING_DATA_FIELD_DEF(uint32_t, attenMaskBatch);
174+TILING_DATA_FIELD_DEF(uint32_t, kvAntiquantSInnerSize);
175+END_TILING_DATA_DEF;
176+REGISTER_TILING_DATA_CLASS(PromptAttentionSingleCoreParamsOp, PromptAttentionSingleCoreParams)
177+ 
178+BEGIN_TILING_DATA_DEF(PromptAttentionSingleCoreTensorSize)
179+TILING_DATA_FIELD_DEF(uint32_t, mmResUbSize);
180+TILING_DATA_FIELD_DEF(uint32_t, pseShiftUbSize);
181+TILING_DATA_FIELD_DEF(uint32_t, attenMaskUbSize);
182+TILING_DATA_FIELD_DEF(uint32_t, maskSize);
183+TILING_DATA_FIELD_DEF(uint32_t, softmaxMaxSize);
184+TILING_DATA_FIELD_DEF(uint32_t, softmaxSumSize);
185+TILING_DATA_FIELD_DEF(uint32_t, softmaxExpSize);
186+TILING_DATA_FIELD_DEF(uint32_t, softmaxValueSize);
187+TILING_DATA_FIELD_DEF(uint32_t, spmTmpSize);
188+TILING_DATA_FIELD_DEF(uint32_t, scmTmpSize);
189+TILING_DATA_FIELD_DEF(uint32_t, bmm2ResUbSize);
190+TILING_DATA_FIELD_DEF(uint32_t, tmpMMResBmm2PreUbSize);
191+TILING_DATA_FIELD_DEF(uint32_t, tmpSoftmaxBmm2UbSize);
192+TILING_DATA_FIELD_DEF(uint32_t, selectSpaceUbSize);
193+TILING_DATA_FIELD_DEF(uint32_t, tmpSoftMaxV2Size);
194+TILING_DATA_FIELD_DEF(uint32_t, mm1TmpUbSize);
195+TILING_DATA_FIELD_DEF(uint32_t, mm2TmpUbSize);
196+TILING_DATA_FIELD_DEF(uint32_t, kvAntiquantUbSize);
197+TILING_DATA_FIELD_DEF(uint32_t, bmm2ResUbMsdSize);
198+TILING_DATA_FIELD_DEF(uint32_t, tempBmm2QueueMsdSize);
199+TILING_DATA_FIELD_DEF(uint32_t, msdInQueueSize);
200+TILING_DATA_FIELD_DEF(uint32_t, msdQRowSumBuffSize);
201+TILING_DATA_FIELD_DEF(uint32_t, msdAMaxTmpBuffSize);
202+TILING_DATA_FIELD_DEF(uint32_t, msdAMaxResBuffSize);
203+TILING_DATA_FIELD_DEF(uint32_t, msdSoftmaxResAmaxBuffSize);
204+TILING_DATA_FIELD_DEF(uint32_t, msdSoftmaxRowSumScaleBuffSize);
205+TILING_DATA_FIELD_DEF(uint32_t, msdScaleBuffSize);
206+TILING_DATA_FIELD_DEF(uint32_t, msdOffsetBuffSize);
207+TILING_DATA_FIELD_DEF(uint32_t, msdTmpMm1BuffSize);
208+TILING_DATA_FIELD_DEF(uint32_t, msdTmpMm2BuffSize);
209+TILING_DATA_FIELD_DEF(uint32_t, msdOutQueueSize);
210+TILING_DATA_FIELD_DEF(uint32_t, msdComputeLines);
211+END_TILING_DATA_DEF;
212+REGISTER_TILING_DATA_CLASS(PromptAttentionSingleCoreTensorSizeOp, PromptAttentionSingleCoreTensorSize)
213+ 
214+BEGIN_TILING_DATA_DEF(PromptAttentionInitOutputParams)
215+TILING_DATA_FIELD_DEF(uint32_t, singleCoreSize);
216+TILING_DATA_FIELD_DEF(int64_t, totalOutputSize);
217+TILING_DATA_FIELD_DEF(int64_t, totalSoftMaxLseOutputSize);
218+TILING_DATA_FIELD_DEF(uint32_t, needInit);
219+TILING_DATA_FIELD_DEF(uint32_t, isOneN);
220+END_TILING_DATA_DEF;
221+REGISTER_TILING_DATA_CLASS(PromptAttentionInitOutputParamsOp, PromptAttentionInitOutputParams)
222+ 
223+BEGIN_TILING_DATA_DEF(PromptFlashAttentionTilingData)
224+TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, bmm1TilingDataRect);
225+TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, bmm2TilingDataRect);
226+ 
227+TILING_DATA_FIELD_DEF_STRUCT(PromptAttentionBaseParams, promptAttentionBaseParams);
228+TILING_DATA_FIELD_DEF_STRUCT(PromptAttentionSeqParams, promptAttentionSeqParams);
229+TILING_DATA_FIELD_DEF_STRUCT(PromptAttentionSingleCoreParams, promptAttentionSingleCoreParams);
230+TILING_DATA_FIELD_DEF_STRUCT(PromptAttentionSingleCoreTensorSize, promptAttentionTensorSizeRect);
231+TILING_DATA_FIELD_DEF_STRUCT(PromptAttentionInitOutputParams, promptAttentionInitOutputParams);
232+ 
233+TILING_DATA_FIELD_DEF_STRUCT(SoftMaxTiling, softmaxTilingDataRect);
234+TILING_DATA_FIELD_DEF_STRUCT(SoftMaxTiling, softmaxFlashTilingDataRect);
235+TILING_DATA_FIELD_DEF_STRUCT(CopyTransposeTiling, transposeTilingDataRect);
236+END_TILING_DATA_DEF;
237+REGISTER_TILING_DATA_CLASS(PromptFlashAttention, PromptFlashAttentionTilingData)
238+ 
239+BEGIN_TILING_DATA_DEF(PFAInputParams)
240+TILING_DATA_FIELD_DEF(int64_t, bSize);
241+TILING_DATA_FIELD_DEF(int64_t, n2Size);
242+TILING_DATA_FIELD_DEF(int64_t, gSize);
243+TILING_DATA_FIELD_DEF(int64_t, s1Size);
244+TILING_DATA_FIELD_DEF(int64_t, s2Size);
245+TILING_DATA_FIELD_DEF(int64_t, alignedS2);
246+TILING_DATA_FIELD_DEF(int64_t, dSize);
247+TILING_DATA_FIELD_DEF(int64_t, valueDSize);
248+TILING_DATA_FIELD_DEF(float, keepProb);
249+TILING_DATA_FIELD_DEF(float, scaleValue);
250+TILING_DATA_FIELD_DEF(int64_t, preTokens);
251+TILING_DATA_FIELD_DEF(int64_t, nextTokens);
252+// in pse encoding scenes, s1 and s2 might not equal with s1, s2 in Q, K
253+TILING_DATA_FIELD_DEF(int64_t, pseS1Size);
254+TILING_DATA_FIELD_DEF(int64_t, pseS2Size);
255+TILING_DATA_FIELD_DEF(uint32_t, pseBSize);
256+TILING_DATA_FIELD_DEF(uint32_t, bandIndex);
257+TILING_DATA_FIELD_DEF(uint32_t, blockSize);
258+TILING_DATA_FIELD_DEF(uint32_t, blockTableDim2);
259+ 
260+// 1: BSH/BSND, 2: SBH, 3: BNSD
261+TILING_DATA_FIELD_DEF(uint8_t, layoutType);
262+// Paged Attention kvcache layout 0: BBH, 1: BNBD, 2: NZ
263+TILING_DATA_FIELD_DEF(uint32_t, paCacheLayoutType);
264+// 0: (B,N2,G,S1,S2), 1: (B,N2,G,1,S2)
265+TILING_DATA_FIELD_DEF(uint8_t, pseShapeType);
266+// 0: (B,N2,G,S1,S2), 1: (B,1,1,S1,S2), 2: (1,1,1,S1,S2)
267+TILING_DATA_FIELD_DEF(uint8_t, attenMaskShapeType);
268+// 0: fp16, 1: bool(uint8)
269+TILING_DATA_FIELD_DEF(uint8_t, attenMaskDataType);
270+// ALL: 0, NONE: 1, ANY: 2, CAUSAL: 3, BAND: 4 };
271+TILING_DATA_FIELD_DEF(uint8_t, attenMaskCompressMode);
272+// 0: high precise, 1: high performance, 2: invalid line high precise
273+TILING_DATA_FIELD_DEF(uint8_t, implMode);
274+TILING_DATA_FIELD_DEF(uint8_t, sparseType);
275+TILING_DATA_FIELD_DEF(uint8_t, fromFused);
276+TILING_DATA_FIELD_DEF(uint8_t, pseEncodeType);
277+TILING_DATA_FIELD_DEF(uint8_t, isSoftMaxLseEnable);
278+TILING_DATA_FIELD_DEF(uint16_t, remain);
279+TILING_DATA_FIELD_DEF(uint32_t, attenMaskS2Size);
280+TILING_DATA_FIELD_DEF(uint32_t, pseType);
281+TILING_DATA_FIELD_DEF(uint32_t, rsv1);
282+TILING_DATA_FIELD_DEF(int64_t, qStartIdx);
283+TILING_DATA_FIELD_DEF(int64_t, kvStartIdx);
284+TILING_DATA_FIELD_DEF(uint32_t, hasLearnableSink);
285+END_TILING_DATA_DEF;
286+REGISTER_TILING_DATA_CLASS(PFAInputParamsOp, PFAInputParams)
287+ 
288+BEGIN_TILING_DATA_DEF(PFAMultiCoreParams)
289+TILING_DATA_FIELD_DEF(int32_t, coreNum);
290+TILING_DATA_FIELD_DEF(int32_t, reserve);
291+// BN2GS1.o
292+TILING_DATA_FIELD_DEF(int64_t, totalSize);
293+// BN2GS1.o / core_num
294+TILING_DATA_FIELD_DEF(int64_t, splitFactorSize);
295+TILING_DATA_FIELD_DEF(int64_t, splitFactorTailSize);
296+TILING_DATA_FIELD_DEF_ARR(int64_t, 48, sparseStartIdx);
297+END_TILING_DATA_DEF;
298+REGISTER_TILING_DATA_CLASS(PFAMultiCoreParamsOp, PFAMultiCoreParams)
299+ 
300+BEGIN_TILING_DATA_DEF(PFACoreParams)
301+TILING_DATA_FIELD_DEF(int32_t, s1BaseSize);
302+TILING_DATA_FIELD_DEF(int32_t, s1BaseTailSize);
303+TILING_DATA_FIELD_DEF(int64_t, s1OuterSize);
304+TILING_DATA_FIELD_DEF(int32_t, s1Vec2BaseSize);
305+TILING_DATA_FIELD_DEF(int32_t, s1Vec2BaseTailSize);
306+TILING_DATA_FIELD_DEF(int64_t, s1Vec2OuterSize);
307+TILING_DATA_FIELD_DEF(int32_t, s2BaseSize);
308+TILING_DATA_FIELD_DEF(int32_t, s2BaseTailSize);
309+TILING_DATA_FIELD_DEF(int64_t, s2OuterSize);
310+TILING_DATA_FIELD_DEF(int32_t, dBaseSize);
311+TILING_DATA_FIELD_DEF(int32_t, dBaseTailSize);
312+TILING_DATA_FIELD_DEF(int64_t, dOuterSize);
313+TILING_DATA_FIELD_DEF(int32_t, bBaseSize);
314+TILING_DATA_FIELD_DEF(int32_t, bBaseTailSize);
315+TILING_DATA_FIELD_DEF(int64_t, bOuterSize);
316+TILING_DATA_FIELD_DEF(int32_t, n2BaseSize);
317+TILING_DATA_FIELD_DEF(int32_t, n2BaseTailSize);
318+TILING_DATA_FIELD_DEF(int64_t, n2OuterSize);
319+TILING_DATA_FIELD_DEF(int32_t, gBaseSize);
320+TILING_DATA_FIELD_DEF(int32_t, gBaseTailSize);
321+TILING_DATA_FIELD_DEF(int64_t, gOuterSize);
322+TILING_DATA_FIELD_DEF(int32_t, nRatio);
323+TILING_DATA_FIELD_DEF(int32_t, rsvd);
324+TILING_DATA_FIELD_DEF(int64_t, s1SparseValidSize);
325+TILING_DATA_FIELD_DEF(int64_t, s2SparseValidSize);
326+TILING_DATA_FIELD_DEF(int64_t, pseAlibiBaseS1);
327+TILING_DATA_FIELD_DEF(int64_t, pseAlibiBaseS2);
328+END_TILING_DATA_DEF;
329+REGISTER_TILING_DATA_CLASS(PFACoreParamsOp, PFACoreParams)
330+ 
331+BEGIN_TILING_DATA_DEF(PFATensorSizeParams)
332+TILING_DATA_FIELD_DEF(int32_t, bmm1ResUbSize);
333+TILING_DATA_FIELD_DEF(int32_t, attenMaskUbSize);
334+TILING_DATA_FIELD_DEF(int32_t, pseUbSize);
335+TILING_DATA_FIELD_DEF(int32_t, dropMaskUbSize);
336+TILING_DATA_FIELD_DEF(int32_t, castUbSize);
337+TILING_DATA_FIELD_DEF(int32_t, softmaxMaxUbSize);
338+TILING_DATA_FIELD_DEF(int32_t, softmaxSumUbSize);
339+TILING_DATA_FIELD_DEF(int32_t, softmaxExpUbSize);
340+TILING_DATA_FIELD_DEF(int32_t, apiTmpBufferBytes);
341+TILING_DATA_FIELD_DEF(int32_t, bmm2ResUbSize);
342+TILING_DATA_FIELD_DEF(int32_t, inputQueBytes);
343+TILING_DATA_FIELD_DEF(int32_t, outputQueBytes);
344+// API buffer use remain space of ub
345+TILING_DATA_FIELD_DEF(int32_t, tmpBufBytes);
346+TILING_DATA_FIELD_DEF(int32_t, softmaxMaxOffsetBytes);
347+TILING_DATA_FIELD_DEF(int32_t, softmaxSumOffsetBytes);
348+TILING_DATA_FIELD_DEF(int32_t, maxSumApiOffsetBytes);
349+TILING_DATA_FIELD_DEF(int32_t, customSoftmaxApiOffsetBytes);
350+TILING_DATA_FIELD_DEF(int32_t, pseTbufOffsetBytes);
351+TILING_DATA_FIELD_DEF(int32_t, dropoutApiOffsetBytes);
352+TILING_DATA_FIELD_DEF(int32_t, maxSumApiSize);
353+TILING_DATA_FIELD_DEF(int32_t, customSoftmaxApiSize);
354+TILING_DATA_FIELD_DEF(int32_t, dropoutApiSize);
355+TILING_DATA_FIELD_DEF(int32_t, attenMaskApiSize);
356+TILING_DATA_FIELD_DEF(int32_t, attenMaskApiOffsetBytes);
357+TILING_DATA_FIELD_DEF(int32_t, bmm1ProcessTInStage2Size);
358+TILING_DATA_FIELD_DEF(int32_t, bmm1ProcessTInStage2OffsetBytes);
359+// workspace
360+TILING_DATA_FIELD_DEF(int32_t, wkspSection1OffsetBytes);
361+TILING_DATA_FIELD_DEF(int32_t, wkspSection2OffsetBytes);
362+END_TILING_DATA_DEF;
363+REGISTER_TILING_DATA_CLASS(PFATensorSizeParamsOp, PFATensorSizeParams)
364+ 
365+BEGIN_TILING_DATA_DEF(MLAGeneralTilingData)
366+TILING_DATA_FIELD_DEF_STRUCT(PFAInputParams, PFAinputParams);
367+TILING_DATA_FIELD_DEF_STRUCT(PFAMultiCoreParams, PFAmultiCoreParams);
368+TILING_DATA_FIELD_DEF_STRUCT(PFACoreParams, PFAcoreParams);
369+TILING_DATA_FIELD_DEF_STRUCT(PFATensorSizeParams, PFAtensorSizeParams);
370+TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, bmm1TilingData);
371+TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, bmm2TilingData);
372+TILING_DATA_FIELD_DEF_STRUCT(SoftMaxTiling, softmaxFlashTilingData);
373+TILING_DATA_FIELD_DEF_STRUCT(CopyTransposeTiling, transposeTilingData);
374+TILING_DATA_FIELD_DEF_STRUCT(CopyTransposeTiling, transposeTilingDataTailCore);
375+END_TILING_DATA_DEF;
376+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_4000000000000000000, MLAGeneralTilingData)
377+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_4000000000000000001, MLAGeneralTilingData)
378+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_4000000000000000002, MLAGeneralTilingData)
379+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_4000000000000000003, MLAGeneralTilingData)
380+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_4000000000000100002, MLAGeneralTilingData)
381+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_4000000000000100003, MLAGeneralTilingData)
382+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_4000000000010000002, MLAGeneralTilingData)
383+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_4000000000010000003, MLAGeneralTilingData)
384+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_4000000000010100002, MLAGeneralTilingData)
385+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_4000000000010100003, MLAGeneralTilingData)
386+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_4000000000020000002, MLAGeneralTilingData)
387+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_4000000000020000003, MLAGeneralTilingData)
388+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_4000000000020100002, MLAGeneralTilingData)
389+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_4000000000020100003, MLAGeneralTilingData)
390+ 
391+BEGIN_TILING_DATA_DEF(PromptFlashAttentionBaseApiTilingData)
392+TILING_DATA_FIELD_DEF_STRUCT(PromptAttentionBaseApiBaseParams, promptAttentionBaseApiBaseParams);
393+TILING_DATA_FIELD_DEF_STRUCT(PromptAttentionSplitCoreParams, promptAttentionSplitCoreParams);
394+END_TILING_DATA_DEF;
395+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_1000000000000112288, PromptFlashAttentionBaseApiTilingData)
396+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_1000000000000122288, PromptFlashAttentionBaseApiTilingData)
397+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_2000000002004000012, PromptFlashAttentionBaseApiTilingData)
398+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_2000000000004001012, PromptFlashAttentionBaseApiTilingData)
399+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_2000000010004001012, PromptFlashAttentionBaseApiTilingData)
400+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_2000000000004000012, PromptFlashAttentionBaseApiTilingData)
401+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_2000000010004000012, PromptFlashAttentionBaseApiTilingData)
402+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_2000000002004010112, PromptFlashAttentionBaseApiTilingData)
403+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_2000000000004010112, PromptFlashAttentionBaseApiTilingData)
404+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_2000000010004010112, PromptFlashAttentionBaseApiTilingData)
405+ 
406+BEGIN_TILING_DATA_DEF(InputParamsRegbase)
407+TILING_DATA_FIELD_DEF(int64_t, bSize);
408+TILING_DATA_FIELD_DEF(int64_t, n2Size);
409+TILING_DATA_FIELD_DEF(int64_t, gSize);
410+TILING_DATA_FIELD_DEF(int64_t, s1Size);
411+TILING_DATA_FIELD_DEF(int64_t, s2Size);
412+TILING_DATA_FIELD_DEF(int64_t, alignedS2);
413+TILING_DATA_FIELD_DEF(int64_t, dSize);
414+TILING_DATA_FIELD_DEF(int64_t, dSizeV);
415+TILING_DATA_FIELD_DEF(int64_t, dSizeRope);
416+TILING_DATA_FIELD_DEF(float, keepProb);
417+TILING_DATA_FIELD_DEF(float, scaleValue);
418+TILING_DATA_FIELD_DEF(int64_t, preTokens);
419+TILING_DATA_FIELD_DEF(int64_t, nextTokens);
420+// in pse encoding scenes, s1 and s2 might not equal with s1, s2 in Q, K
421+TILING_DATA_FIELD_DEF(int64_t, pseS1Size);
422+TILING_DATA_FIELD_DEF(int64_t, pseS2Size);
423+TILING_DATA_FIELD_DEF(uint32_t, pseBSize);
424+TILING_DATA_FIELD_DEF(uint32_t, bandIndex);
425+ 
426+// 1: BSH/BSND, 2: SBH, 3: BNSD
427+TILING_DATA_FIELD_DEF(uint8_t, layoutType);
428+// 0: (B,N2,G,S1,S2), 1: (B,N2,G,1,S2)
429+TILING_DATA_FIELD_DEF(uint8_t, pseShapeType);
430+// 0: (B,N2,G,S1,S2), 1: (B,1,1,S1,S2), 2: (1,1,1,S1,S2)
431+TILING_DATA_FIELD_DEF(uint8_t, attenMaskShapeType);
432+// 0: fp16, 1: bool(uint8)
433+TILING_DATA_FIELD_DEF(uint8_t, attenMaskDataType);
434+// ALL: 0, NONE: 1, ANY: 2, CAUSAL: 3, BAND: 4 };
435+TILING_DATA_FIELD_DEF(uint8_t, attenMaskCompressMode);
436+// 0: high precise, 1: high performance, 2: invalid line high precise
437+TILING_DATA_FIELD_DEF(uint8_t, implMode);
438+TILING_DATA_FIELD_DEF(uint8_t, sparseType);
439+TILING_DATA_FIELD_DEF(uint8_t, needDropMaskOp);
440+TILING_DATA_FIELD_DEF(uint8_t, dropMaskOuter);
441+TILING_DATA_FIELD_DEF(uint8_t, pseEncodeType);
442+TILING_DATA_FIELD_DEF(uint16_t, remain);
443+TILING_DATA_FIELD_DEF(uint32_t, attenMaskS2Size);
444+TILING_DATA_FIELD_DEF(uint32_t, pseType);
445+TILING_DATA_FIELD_DEF(uint32_t, rsv1);
446+TILING_DATA_FIELD_DEF(int64_t, qStartIdx);
447+TILING_DATA_FIELD_DEF(int64_t, kvStartIdx);
448+TILING_DATA_FIELD_DEF(int64_t, s1SparseValidSize);
449+TILING_DATA_FIELD_DEF(int64_t, s2SparseValidSize);
450+TILING_DATA_FIELD_DEF(int64_t, seed);
451+TILING_DATA_FIELD_DEF(int64_t, offset);
452+TILING_DATA_FIELD_DEF(int64_t, keepProbUint8);
453+TILING_DATA_FIELD_DEF(int64_t, pseAlibiBaseS1);
454+TILING_DATA_FIELD_DEF(int64_t, pseAlibiBaseS2);
455+ 
456+// PFA
457+TILING_DATA_FIELD_DEF(uint8_t, deqScaleFlag);
458+TILING_DATA_FIELD_DEF(uint8_t, deqScale2Flag);
459+TILING_DATA_FIELD_DEF(uint8_t, isActualSeqLengthsNull);
460+TILING_DATA_FIELD_DEF(uint8_t, isActualSeqLengthsKVNull);
461+TILING_DATA_FIELD_DEF(uint32_t, actualSeqLengthsSize);
462+TILING_DATA_FIELD_DEF(uint32_t, actualSeqLengthsKVSize);
463+TILING_DATA_FIELD_DEF(uint8_t, isKvContinuous);
464+TILING_DATA_FIELD_DEF(uint8_t, fromFused);
465+TILING_DATA_FIELD_DEF(uint8_t, isBSNDOut);
466+TILING_DATA_FIELD_DEF(uint8_t, isGqa);
467+TILING_DATA_FIELD_DEF(uint8_t, isSoftMaxLseEnable);
468+TILING_DATA_FIELD_DEF(uint8_t, isActualSharedPrefixLenNull);
469+TILING_DATA_FIELD_DEF(uint8_t, isQHasLeftPadding);
470+TILING_DATA_FIELD_DEF(uint8_t, isKVHasLeftPadding);
471+TILING_DATA_FIELD_DEF(uint32_t, ropeHeadSize);
472+TILING_DATA_FIELD_DEF(uint32_t, prefixSeqInnerSize);
473+TILING_DATA_FIELD_DEF(uint32_t, headNumRatio);
474+TILING_DATA_FIELD_DEF(int32_t, blockSize);
475+TILING_DATA_FIELD_DEF(int32_t, blockTableDim2);
476+TILING_DATA_FIELD_DEF(int32_t, paBlockNumSum);
477+TILING_DATA_FIELD_DEF(uint32_t, attenMaskS1Size);
478+TILING_DATA_FIELD_DEF(uint32_t, kvSplitPart);
479+TILING_DATA_FIELD_DEF(uint32_t, accumOutSize);
480+TILING_DATA_FIELD_DEF(uint32_t, logSumExpSize);
481+ 
482+TILING_DATA_FIELD_DEF(uint8_t, paLayoutType);
483+TILING_DATA_FIELD_DEF(uint8_t, isRowInvalid);
484+TILING_DATA_FIELD_DEF(uint8_t, isPostQuantPerChnl);
485+TILING_DATA_FIELD_DEF(uint8_t, isPostQuantBF16);
486+TILING_DATA_FIELD_DEF(uint16_t, antiquantPerTensorFlag);
487+TILING_DATA_FIELD_DEF(uint16_t, antiquantPerHeadFlag);
488+TILING_DATA_FIELD_DEF(uint32_t, antiquantParaSeqSize);
489+END_TILING_DATA_DEF;
490+REGISTER_TILING_DATA_CLASS(InputParamsRegbaseOp, InputParamsRegbase)
491+ 
492+BEGIN_TILING_DATA_DEF(MultiCoreParamsRegbase)
493+TILING_DATA_FIELD_DEF(int32_t, coreNum);
494+TILING_DATA_FIELD_DEF(int64_t, totalSize);
495+TILING_DATA_FIELD_DEF(int64_t, s1OuterSize);
496+TILING_DATA_FIELD_DEF(int64_t, splitFactorSize);
497+TILING_DATA_FIELD_DEF(int64_t, splitFactorTailSize);
498+TILING_DATA_FIELD_DEF_ARR(uint32_t, 48, bnStartIdx);
499+TILING_DATA_FIELD_DEF_ARR(int64_t, 48, sparseStartIdx);
500+END_TILING_DATA_DEF;
501+REGISTER_TILING_DATA_CLASS(MultiCoreParamsRegbaseOp, MultiCoreParamsRegbase)
502+ 
503+BEGIN_TILING_DATA_DEF(DropmaskParamsRegbase)
504+TILING_DATA_FIELD_DEF(int32_t, multiCoreFactorSize);
505+TILING_DATA_FIELD_DEF(int32_t, baseUbCalSize);
506+TILING_DATA_FIELD_DEF(int64_t, multiCoreTotalSize);
507+TILING_DATA_FIELD_DEF(int64_t, shapeTotalSize);
508+END_TILING_DATA_DEF;
509+REGISTER_TILING_DATA_CLASS(DropmaskParamsRegbaseOp, DropmaskParamsRegbase)
510+ 
511+BEGIN_TILING_DATA_DEF(InitOutputParams)
512+TILING_DATA_FIELD_DEF(uint32_t, singleCoreSize);
513+TILING_DATA_FIELD_DEF(uint8_t, needInit);
514+TILING_DATA_FIELD_DEF(uint8_t, isOneN);
515+TILING_DATA_FIELD_DEF_ARR(uint8_t, 2, rsvd);
516+TILING_DATA_FIELD_DEF(int64_t, totalOutputSize);
517+TILING_DATA_FIELD_DEF(int64_t, totalSoftMaxLseOutputSize);
518+END_TILING_DATA_DEF;
519+REGISTER_TILING_DATA_CLASS(InitOutputParamsOp, InitOutputParams)
520+ 
521+BEGIN_TILING_DATA_DEF(FlashAttentionScoreSimplifiedTilingData)
522+TILING_DATA_FIELD_DEF_STRUCT(InputParamsRegbase, inputParamsRegbase);
523+TILING_DATA_FIELD_DEF_STRUCT(MultiCoreParamsRegbase, multiCoreParamsRegbase);
524+TILING_DATA_FIELD_DEF_STRUCT(DropmaskParamsRegbase, dropmaskParamsRegbase);
525+TILING_DATA_FIELD_DEF_STRUCT(InitOutputParams, initOutputParams);
526+END_TILING_DATA_DEF;
527+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_1000000000000000090, FlashAttentionScoreSimplifiedTilingData)
528+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_10000000000000090, FlashAttentionScoreSimplifiedTilingData)
529+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_1002312000040001212, FlashAttentionScoreSimplifiedTilingData)
530+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_1002312000040021212, FlashAttentionScoreSimplifiedTilingData)
531+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_1001311000000001212, FlashAttentionScoreSimplifiedTilingData)
532+REGISTER_TILING_DATA_CLASS(PromptFlashAttention_1001311000000021212, FlashAttentionScoreSimplifiedTilingData)
533+ 
534+class BufferNum {
535+ public:
536+ // sum and max always use fp32, shape is (S1, 1), inner axis align 32B.
537+ size_t bufferS1S2Num; // unit: input dtype
538+ size_t bufferS1DNum;
539+ size_t bufferExpNum; // unit: input dtype, shape: [S1, 1], inner axis align 32B.
540+};
541+ 
542+class PromptFlashAttentionTiling : public FiaTilingBase {
543+ public:
544+ explicit PromptFlashAttentionTiling(gert::TilingContext *context)
545+ : FiaTilingBase(context), ascendcPlatform(nullptr) {}
546+ ~PromptFlashAttentionTiling() override = default;
547+ ge::graphStatus RunBigKernelTilingWithParams(ContextParamsForPFATiling &contextKeyParams, uint64_t &tilingKey,
548+ uint32_t &numBlocksToBeSet, PromptFlashAttentionTilingData &tilingData);
549+ ge::graphStatus PromptFlashAttentionSetTilingData(
550+ gert::TilingContext *context, PromptFlashAttentionTilingData &tilingData);
551+ bool CheckNonEmptyShapeExceptions(
552+ ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *shape, const std::string &sName);
553+ bool fromPFA_ = true;
554+ bool CheckBaseApiNonEmptyShapeExceptions(
555+ ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *shape, const std::string &sName);
556+ 
557+ protected:
558+ void InitTilingInfo(TilingInfo *tilingInfo) override {}
559+ bool IsCapable() override { return true; }
560+ ge::graphStatus DoOpTiling() override;
561+ ge::graphStatus ConvertContextToPFAParams(
562+ gert::TilingContext *context, ContextParamsForPFATiling &contextKeyParams) const;
563+ ge::graphStatus TilingGetTilingKeyAttentionAscendC(uint64_t &tilingKey, ContextParamsForPFATiling &contextKeyParams,
564+ bool useNewTiling, PromptFlashAttentionTilingData &tilingData);
565+ void PromptFlashAttentionSplitNS(const ContextParamsForPFATiling &contextKeyParams,
566+ PromptFlashAttentionTilingData &tilingData, uint32_t curCoreNum, const std::vector<int64_t> &actualSeqLengths);
567+ void PromptFlashAttentionSplitNSNew(const ContextParamsForPFATiling &contextKeyParams,
568+ PromptFlashAttentionTilingData &tilingData, uint32_t curCoreNum, const std::vector<int64_t> &actualSeqLengths,
569+ const std::vector<int64_t> &actualSeqLengthsKV, int64_t actualSharedPrefixLen, bool useBalanceTiling);
570+ void GetPreNextTokensLeftUp(PromptFlashAttentionTilingData &tilingData, uint32_t actualSeqLength,
571+ uint32_t actualSeqLengthKV, int64_t &preTokensLeftUp, int64_t &nextTokensLeftUp);
572+ void SetSplitCoreMode(PromptFlashAttentionTilingData &tilingData, uint32_t sOuterFactor);
573+ void PromptFlashAttentionSplitSeqOneN(
574+ PromptFlashAttentionTilingData &tilingData, uint32_t curCoreNum, bool isVectorCore);
575+ bool EnableMTE2BmmPipe(PromptFlashAttentionTilingData &tilingData, matmul_tiling::MatmulApiTiling &bmm,
576+ TCubeTiling &bmmTilingData, uint32_t sOuterFactor, uint32_t sInnerFactor);
577+ void EnableBmmDoubleBuffer(TCubeTiling &bmmTilingData);
578+ void PromptFlashAttention310PSetBmm1(matmul_tiling::MatmulApiTiling &bmm1);
579+ void PromptFlashAttention310PSetBmm2(matmul_tiling::MatmulApiTiling &bmm2);
580+ bool PromptFlashAttentionCheckBmm1(PromptFlashAttentionTilingData &tilingData, TCubeTiling &bmm1TilingData,
581+ int64_t l1SizeRemain, int64_t l0CSize, uint32_t sOuterFactor, uint32_t sInnerFactor, bool allGM = false,
582+ bool autoBaseMNK = false);
583+ bool PromptFlashAttentionCheckBmm2(PromptFlashAttentionTilingData &tilingData, TCubeTiling &bmm1TilingData,
584+ int64_t l1SizeRemain, int64_t l0CSize, uint32_t sOuterFactor, uint32_t sInnerFactor, uint32_t dSplitFactor,
585+ bool allGM = false, bool autoBaseMNK = false);
586+ void PromptFlashAttentionSetTensorSize(PromptFlashAttentionTilingData &tilingData,
587+ PromptAttentionSingleCoreTensorSize &tensorSize, uint32_t sOuterFactor, uint32_t sInnerFactor);
588+ bool PromptFlashAttentionCheckArgsLegal(PromptFlashAttentionTilingData &tilingData, int64_t ubSize, int64_t l1Size,
589+ int64_t l0CSize, uint32_t typeByteSize, uint32_t &sOuterFactor, uint32_t sInnerFactor, bool &updateDiv,
590+ uint32_t maskTypeSize, uint32_t dSplitFactor);
591+ ge::graphStatus AdjustBasicBlock(PromptFlashAttentionTilingData &tilingData, uint32_t &sOuterFactor);
592+ ge::graphStatus PromptFlashAttentionApiTiling(PromptFlashAttentionTilingData &tilingData, uint32_t typeSize,
593+ uint32_t sOuterFactor, uint32_t softmaxSInnerFactor, uint32_t softmaxSOuterFactor);
594+ ge::graphStatus GetRectangleFactor(uint32_t seqSplit, std::queue<uint32_t> &sQueue, int32_t threshold = 16);
595+ ge::graphStatus SetInputLayout(const char *layout);
596+ bool GetApiTmpSize(const uint32_t sOuterFactor, const uint32_t sInnerFactor, const uint32_t typeByteSize);
597+ uint32_t CalculateL1SizeUsed(PromptFlashAttentionTilingData &tilingData, const uint32_t typeByteSize);
598+ bool CheckInputDimAndHeadNum(ContextParamsForPFATiling &contextKeyParams, uint32_t nQAttr, uint32_t nKVAttr);
599+ bool SetTilingHeadNumRatio(ContextParamsForPFATiling &contextKeyParams, const int64_t *numQueryHeads,
600+ const int64_t *numKeyValueHeads, PromptFlashAttentionTilingData &tilingData);
601+ void PromptFlashAttentionInitOutputSplit(
602+ uint64_t totalSize, PromptFlashAttentionTilingData &tilingData, uint32_t curCoreNum);
603+ void PromptFlashAttentionInitSoftmaxLseOutputSplit(uint64_t totalSize, PromptFlashAttentionTilingData &tilingData);
604+ void Align(uint32_t &num);
605+ ge::graphStatus GetBasicShape(uint32_t &b, uint32_t &s, uint32_t &h, uint32_t &seqInnerSize,
606+ const gert::StorageShape *queryShape, const gert::StorageShape *keyShape, const uint32_t n);
607+ ge::graphStatus GetBasicShape310P(uint32_t &b, uint32_t &bKV, uint32_t &s, uint64_t &h, uint32_t &seqInnerSize,
608+ const gert::StorageShape *queryShape, const gert::StorageShape *keyShape, const uint64_t n,
609+ size_t actualLenDims, size_t actualLenDimsKV);
610+ ge::graphStatus GetBasicShape910B(uint32_t &b, uint32_t &s, uint32_t &h, uint32_t &seqInnerSize,
611+ const gert::StorageShape *queryShape, const gert::StorageShape *keyShape, const uint32_t n);
612+ size_t GetPFAWorkSpaceSize(PromptFlashAttentionTilingData &tilingData);
613+ void GetMatMulType(matmul_tiling::DataType &mmInputType, matmul_tiling::DataType &mmOutputType);
614+ ge::graphStatus CheckKeyValueParamsConsistency(const ContextParamsForPFATiling &contextKeyParams);
615+ bool CheckActualSeqLength(ContextParamsForPFATiling &contextKeyParams, uint32_t b, uint32_t sQ, uint32_t sKV,
616+ const gert::Tensor *actualSeqLenQ, const gert::Tensor *actualSeqLenKV, InputLayout inLayout,
617+ PromptFlashAttentionTilingData &tilingData);
618+ bool CheckPseShiftTypeAndShape(ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *pseShiftShape,
619+ uint32_t b, uint64_t n, uint32_t s1, uint32_t s2);
620+ ge::graphStatus processPageAttentionInputFlag(ContextParamsForPFATiling &contextKeyParams);
621+ bool checkPAKeyValueDimsWhenBBH(ContextParamsForPFATiling &contextKeyParams, int32_t keyDim1, int32_t keyDim2,
622+ int32_t keyDim3, int64_t blockNumValid, const int32_t *curBlockSize, int64_t h, int64_t headNumRatio);
623+ bool checkPAKeyValueDimsWhenBNBD(ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *keyShape,
624+ const gert::StorageShape *valueShape, int32_t keyDim1, int32_t keyDim2, int32_t keyDim3, int64_t blockNumValid,
625+ const int32_t *curBlockSize, int64_t n, int64_t h, int64_t headNumRatio);
626+ bool checkPAKeyValueDimsWhenNZ(ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *keyShape,
627+ const gert::StorageShape *valueShape, int32_t keyDim1, int32_t keyDim2, int32_t keyDim3, int64_t blockNumValid,
628+ const int32_t *curBlockSize, int64_t h, int64_t n, int64_t headNumRatio);
629+ bool checkPABlockSizeAndBlockTable(ContextParamsForPFATiling &contextKeyParams, const gert::Tensor *actualSeqLenKV,
630+ const int32_t *curBlockSize, int64_t b);
631+ bool culActSeqLenParamsWhenPA(ContextParamsForPFATiling &contextKeyParams, const gert::Tensor *actualSeqLenKV,
632+ int64_t &blockNumValid, int32_t &maxBlockNumPerBatch, int32_t tempBlockSize, const int32_t *curBlockSize,
633+ int64_t b);
634+ bool CheckPAKeyValueParams(ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *keyShape,
635+ const gert::StorageShape *valueShape, int64_t blockNumValid, const int32_t *curBlockSize, int64_t n, int64_t h,
636+ int64_t headNumRatio);
637+ bool CheckPASparseMode(ContextParamsForPFATiling &contextKeyParams);
638+ bool CheckPAWhenBaseApi(ContextParamsForPFATiling &contextKeyParams, const gert::Tensor *actualSeqLenQ,
639+ const gert::Tensor *actualSeqLenKV, int64_t n, int64_t h, int64_t headNumRatio);
640+ bool CheckPATypeAndShape(ContextParamsForPFATiling &contextKeyParams, const gert::Tensor *actualSeqLenKV, int32_t b,
641+ int64_t n, int64_t h, int64_t headNumRatio);
642+ bool CheckAttenMaskShape(ContextParamsForPFATiling &contextKeyParams, const int32_t *sparseMode,
643+ const gert::StorageShape *attenMaskShape, const uint64_t sQ, const uint64_t sK, const uint32_t batchSize);
644+ bool CheckPAAntiquantSupportScenarios(
645+ ContextParamsForPFATiling &contextKeyParams, PromptFlashAttentionTilingData &tilingData);
646+ bool CheckPerchannelAntiquantParamsShape(ContextParamsForPFATiling &contextKeyParams,
647+ const gert::StorageShape *antiquantScaleShape, const gert::StorageShape *antiquantOffsetShape, const uint64_t n,
648+ const uint32_t d, const uint64_t h, uint32_t paramFirstDim) const;
649+ bool CheckPerchannelBSNDParamsShape(ContextParamsForPFATiling &contextKeyParams,
650+ const gert::StorageShape *antiquantScaleShape, const gert::StorageShape *antiquantOffsetShape, const uint64_t n,
651+ const uint32_t d, uint32_t paramFirstDim) const;
652+ bool CheckAntiquantParamsShape(ContextParamsForPFATiling &contextKeyParams,
653+ const gert::StorageShape *antiquantScaleShape, const gert::StorageShape *antiquantOffsetShape, const uint64_t n,
654+ const uint32_t d, const uint64_t h, PromptFlashAttentionTilingData &tilingData);
655+ ge::graphStatus CheckPostQuantParams(
656+ const ContextParamsForPFATiling &contextKeyParams, uint64_t h, uint64_t n) const;
657+ ge::graphStatus PromptFlashAttentionCVDiffSetTensorSize(PromptFlashAttentionTilingData &tilingData,
658+ PromptAttentionSingleCoreTensorSize &tensorSize, uint32_t sOuterFactor, uint32_t sInnerFactor,
659+ uint32_t softmaxSOuterFactor);
660+ bool PromptFlashAttentionComputeCVDiffParams(PromptFlashAttentionTilingData &tilingData, int64_t ubSize,
661+ int64_t l1Size, int64_t l0CSize, uint32_t typeByteSize, uint32_t &sOuterFactor, uint32_t &sInnerFactor,
662+ uint32_t maskTypeSize, uint32_t &softmaxSOuterFactor);
663+ bool FindOptimalTilingBasicBLock(PromptFlashAttentionTilingData &tilingData, uint32_t &sOuterFactor,
664+ uint32_t &sInnerFactor, uint32_t &softmaxSOuterFactor, int64_t ubSize, uint32_t typeByteSize,
665+ uint32_t maskTypeSize);
666+ bool FindOptimalTilingSouter(PromptFlashAttentionTilingData &tilingData, uint32_t &sOuterFactor,
667+ uint32_t &sInnerFactor, uint32_t &softmaxSOuterFactor, int64_t ubSize, uint32_t typeByteSize,
668+ uint32_t maskTypeSize);
669+ void InferTilingMod(const ContextParamsForPFATiling &contextKeyParams, const std::vector<int64_t> &actualSeqLengths,
670+ const std::vector<int64_t> &actualSeqLengthsKV, uint32_t actualSeqArrayLen, uint32_t hDivN,
671+ uint32_t seqInnerSize, int32_t sparseModeVal);
672+ ge::graphStatus AdjustCVTiling(uint64_t hDivN, uint64_t n, int64_t middleActualSeqLengths, int64_t ubSize,
673+ int64_t l1Size, int64_t l0CSize, uint32_t maskElemSize, uint32_t &sOuterFactor, uint32_t &sInnerFactor,
674+ PromptFlashAttentionTilingData &tilingData);
675+ ge::graphStatus AdjustCVTilingCVDiff(int64_t ubSize, int64_t l1Size, int64_t l0CSize, uint32_t maskElemSize,
676+ uint32_t &sOuterFactor, uint32_t &sInnerFactor, uint32_t &softmaxSOuterFactor,
677+ PromptFlashAttentionTilingData &tilingData);
678+ bool CheckSparseModeRightDown(ContextParamsForPFATiling &contextKeyParams,
679+ const std::vector<int64_t> &actualSeqLengths, const std::vector<int64_t> &actualSeqLengthsKV, size_t lenDims);
680+ ge::graphStatus GetAndCheckEmptyQueryShape(
681+ ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *queryShape) const;
682+ void UpdateTilingKeyFlag(const ContextParamsForPFATiling &contextKeyParams, uint64_t &tilingKey);
683+ int64_t PromptFlashAttentionSetMsdUbSize(PromptFlashAttentionTilingData &tilingData,
684+ PromptAttentionSingleCoreTensorSize &tensorSize, int32_t sInnerFactorTmp) const;
685+ ge::graphStatus CheckIOType(ContextParamsForPFATiling &contextKeyParams, PromptFlashAttentionTilingData &tilingData,
686+ int32_t &outputDataTypeSize);
687+ ge::graphStatus CheckD(ContextParamsForPFATiling &contextKeyParams) const;
688+ ge::graphStatus CheckDimNums(ContextParamsForPFATiling &contextKeyParams);
689+ ge::graphStatus CheckMaskType(ContextParamsForPFATiling &contextKeyParams,
690+ PromptFlashAttentionTilingData &tilingData, uint32_t &maskElemSize) const;
691+ void SetMaskSize(const gert::StorageShape *attenMaskShape, PromptFlashAttentionTilingData &tilingData) const;
692+ ge::graphStatus CheckShape(ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *queryShape,
693+ const gert::StorageShape *keyShape, const gert::StorageShape *valueShape, const gert::StorageShape *outShape,
694+ const gert::StorageShape *pseShiftShape, const gert::StorageShape *attenMaskShape);
695+ ge::graphStatus CheckBaseAPISupportScenarios(ContextParamsForPFATiling &contextKeyParams);
696+ size_t GetPFABaseApiWorkSpaceSize(const uint32_t &numBlocksToBeSet);
697+ ge::graphStatus TilingGetBaseApiTilingKeyAttentionAscendC(
698+ uint64_t &tilingKey, ContextParamsForPFATiling &contextKeyParams);
699+ ge::graphStatus CheckBaseApiRequiredInput(ContextParamsForPFATiling &contextKeyParams);
700+ ge::graphStatus CheckBaseApiOptionalInput(ContextParamsForPFATiling &contextKeyParams);
701+ ge::graphStatus CheckBaseApiPse(
702+ ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *pseShiftShape);
703+ void SetBaseApiTilingData(ContextParamsForPFATiling &contextKeyParams, std::vector<int64_t> &actualSeqLengths,
704+ std::vector<int64_t> &actualSeqLengthsKV);
705+ void SetBaseApiSeqTilingData(ContextParamsForPFATiling &contextKeyParams, std::vector<int64_t> &actualSeqLengths,
706+ std::vector<int64_t> &actualSeqLengthsKV);
707+ 
708+ ge::graphStatus CheckBaseApiMaskBasic(ContextParamsForPFATiling &contextKeyParams,
709+ const gert::StorageShape *pseShiftShape, bool isLongSeq, uint32_t batchSize) const;
710+ ge::graphStatus CheckBaseApiMaskVal(ContextParamsForPFATiling &contextKeyParams,
711+ const gert::StorageShape *pseShiftShape, const std::pair<std::vector<int64_t>, std::string> maskShape) const;
712+ ge::graphStatus CheckBaseApiAlibiMask(ContextParamsForPFATiling &contextKeyParams,
713+ const gert::StorageShape *pseShiftShape, uint32_t batchSize, int32_t maxSeqLen, int32_t maxKvSeqLen,
714+ uint32_t kvHead, bool compressHead) const;
715+ ge::graphStatus CheckBaseApiNormMask(ContextParamsForPFATiling &contextKeyParams,
716+ const gert::StorageShape *pseShiftShape, int32_t maskType, uint32_t batchSize, int32_t maxSeqLen,
717+ int32_t maxKvSeqLen, bool compressHead) const;
718+ 
719+ ge::graphStatus SetBaseApiPseInfo(
720+ ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *pseShiftShape);
721+ void SetBaseApiOtherMaskInfo(ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *pseShiftShape);
722+ ge::graphStatus SetBaseApiAlibiMaskInfo(
723+ ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *pseShiftShape);
724+ ge::graphStatus AtbSplitBlock(ContextParamsForPFATiling &contextKeyParams);
725+ uint32_t CalcTschNumBlocks(uint32_t sliceNum, uint32_t aicCoreNum, uint32_t aivCoreNum) const;
726+ void CalcUBSize();
727+ void SetDataCopyTransposeTiling();
728+ void SetSoftMaxTiling();
729+ bool SetBmm1TilingInput(int64_t tmpS1BasicBlock, int64_t tmpS2BasicBlock, matmul_tiling::MatmulApiTiling &bmm1);
730+ bool SetBmm2TilingInput(int64_t tmpS1BasicBlock, int64_t tmpS2BasicBlock, matmul_tiling::MatmulApiTiling &bmm2);
731+ bool SetMatMulTiling(int64_t tmpS1BasicBlock, int64_t tmpS2BasicBlock, matmul_tiling::MatmulApiTiling &bmm1,
732+ matmul_tiling::MatmulApiTiling &bmm2);
733+ bool SetMatMulTiling(int64_t tmpS1BasicBlock, int64_t tmpS2BasicBlock);
734+ void CalcS1S2BasicBlock(const BufferNum &bufferNum);
735+ int64_t CalcMaxS1BasicBlockSize(int64_t actualD, const BufferNum &bufferNum);
736+ int64_t CalcMaxS2BasicBlockSize(const BufferNum &bufferNum, int64_t tmpS1BasicBlock);
737+ bool IsBasicBlockInSoftMax(const ge::Shape &shape) const;
738+ void GetBufferNum(BufferNum &bufferNum) const;
739+ void SetMultiBatchCoreParams();
740+ void MatchTemplate(uint32_t valueD);
741+ void SetTensorSizeParams();
742+ bool SetSparseStartIdx(const std::vector<int64_t> &sparseValidArray, PFAMultiCoreParams &multiCoreParams);
743+ 
744+ // TND新增
745+ bool InputLayoutIsTNDLike() const;
746+ int64_t GetTFromInputShape(uint32_t inputIdx, const gert::StorageShape *shape) const;
747+ int64_t GetNFromInputShape(uint32_t inputIdx, const gert::StorageShape *shape) const;
748+ int64_t GetDFromInputShape(uint32_t inputIdx, const gert::StorageShape *shape) const;
749+ int64_t GetTFromOutputShape(const gert::StorageShape *shape) const;
750+ int64_t GetNFromOutputShape(const gert::StorageShape *shape) const;
751+ void GetActualSeqLenData(int64_t inputIdx, std::array<int64_t, MAX_VAR_LEN_SEQ_LEN> &res, int64_t &actualLen) const;
752+ void SetMultiCoreParamsTND();
753+ void SetSparseParamsTND();
754+ void InitSparseValidArrayTND(std::vector<int64_t> &sparseValidArray);
755+ bool SetSparseStartIdxTND(const std::vector<int64_t> &sparseValidArray, PFAMultiCoreParams &multiCoreParams) const;
756+ int64_t GetS2RealSize(uint8_t sparseType, int32_t bOutIdx, int64_t s1OutIdx);
757+ bool BalanceLoad(const std::vector<int64_t> &sparseValidArray, PFAMultiCoreParams &multiCoreParams,
758+ std::vector<int64_t> &localValue, std::vector<int64_t> &sparseStartIdx) const;
759+ void InitLoadValue(const std::vector<int64_t> &sparseValidArray, int64_t validAivNum, int64_t totalSize,
760+ const std::vector<int64_t> &sparseStartIdx, std::vector<int64_t> &localValue) const;
761+ ge::graphStatus CheckInputShapeWhenLayoutIsTND(ContextParamsForPFATiling &contextKeyParams) const;
762+ ge::graphStatus CheckActSeqWhenLayoutIsTND(ContextParamsForPFATiling &contextKeyParams) const;
763+ 
764+ ge::graphStatus CheckVarLenPreNextToken(ContextParamsForPFATiling &contextKeyParams, int32_t sparseMode,
765+ int64_t sparsePreTokens, int64_t sparseNextTokens);
766+ ge::graphStatus CheckLearnableSinkWhenLayoutIsTND(ContextParamsForPFATiling &contextKeyParams) const;
767+ 
768+ protected:
769+ ContextParamsForPFATiling *contextKeyParamsPtr = nullptr;
770+ int64_t ubSizeRemain = 1;
771+ bool isSOuterNoTail = true;
772+ bool isSInnerNoTail = true;
773+ bool isDNoTail = true;
774+ bool enableKvAntiquant = false;
775+ bool enableMsd = false;
776+ bool enableQuantBF16 = false;
777+ bool enableMatmulNorm = false;
778+ bool enablePA = false;
779+ bool isKVHasPrefix = false;
780+ InputLayout inputLayout = InputLayout::BSH;
781+ InputLayout inputKvLayout = InputLayout::BSH;
782+ ge::DataType inputType{ge::DT_FLOAT16};
783+ ge::DataType outputType{ge::DT_FLOAT16};
784+ ge::DataType intputKeyType{ge::DT_FLOAT16};
785+ ge::DataType intputValueType{ge::DT_FLOAT16};
786+ ge::DataType pseShiftElemType{ge::DT_FLOAT16};
787+ uint32_t dataTypeSize = FLOAT32SIZE;
788+ uint32_t coreNum = 0;
789+ uint32_t aivNum = 0;
790+ uint32_t aicNum = 0;
791+ uint32_t typeByteNum = 0;
792+ uint32_t outputTypeByteNum = 0;
793+ uint32_t softmaxTypeByteNum = 0;
794+ uint32_t pseShiftTypeByteNum = 0;
795+ uint32_t pseShiftElemSize = 0;
796+ uint32_t pseMaskMaxSize = 0;
797+ uint32_t pseShiftBatch = 0;
798+ uint32_t pseShiftS1 = 0;
799+ uint32_t pseShiftS2 = 0;
800+ uint32_t usePseShift = 0;
801+ uint32_t tmpS2 =
802+ 0; // In the PA scenario, there is no S2 axis. Use the change amount to normalize the S2 length in both PA and non PA scenarios
803+ int32_t blockSize = 128;
804+ int32_t blockTableDim2 = 1;
805+ int32_t PABlockNumSum = 1;
806+ uint32_t maskTypeByteNum;
807+ uint32_t maxQuerySeq = 0;
808+ int64_t apiTmpSize = 1;
809+ uint32_t softmaxDataTypeNZ_ = FLOAT32SIZE;
810+ uint32_t softmaxDataTypeSize = FLOAT32SIZE; // BF16 calculates through FP32
811+ platform_ascendc::SocVersion curShortSocName;
812+ uint32_t dataTypeSize_ = 4;
813+ uint32_t layoutType = 0;
814+ uint32_t PAlayoutType = 0;
815+ platform_ascendc::PlatformAscendC ascendcPlatform;
816+ TilingMod tilingMod = TilingMod::CVSAME;
817+ SplitCoreMode splitCoreMode = SplitCoreMode::SPLIT_NBS_VECTOR;
818+ uint32_t splitD = 0;
819+ uint32_t splitS2 = 1; // It can only be 0 when the D axis is split
820+ uint64_t innerPrecise = HIGH_PERFORMANCE;
821+ size_t defaultSysWorkspaceSize;
822+ matmul_tiling::PlatformInfo ascendPlatformInfo;
823+ 
824+ int64_t alignedS1 = 0;
825+ int64_t alignedS2 = 0;
826+ int64_t alignedD = 0;
827+ 
828+ int64_t s1BasicBlock = 0;
829+ int64_t s2BasicBlock = 0;
830+ int64_t s1BasicBlockBest = 0;
831+ int64_t s1VecBasicBlock = 0;
832+ int64_t dBasicBlock = 0;
833+ int64_t batchBasic = 1LL;
834+ int64_t nRatio = 0;
835+ 
836+ int64_t s1Size = 0;
837+ int64_t s2Size = 0;
838+ int64_t dSize = 0;
839+ int64_t valueDSize = 0;
840+ int64_t s1SparseValidSize = 0;
841+ int64_t s2SparseValidSize = 0;
842+ 
843+ int64_t apiMaxUBSize = 0;
844+ 
845+ bool atbRunFlag_ = false;
846+ bool mlaRunFlag_ = false;
847+ 
848+ // TND新增
849+ int64_t realT1Size = 0;
850+ int64_t realT2Size = 0;
851+ std::array<int64_t, MAX_VAR_LEN_SEQ_LEN> actualSeqLenData;
852+ std::array<int64_t, MAX_VAR_LEN_SEQ_LEN> actualSeqLenKvData;
853+ int64_t accumS1 = 0;
854+ int64_t accumS2 = 0;
855+ int64_t bandIndex = 0;
856+ 
857+ int64_t bSize = 0;
858+ int64_t gSize = 0;
859+ int64_t n1Size = 0;
860+ int64_t n2Size = 0;
861+ int64_t s1StrideSize = 0; // query Shape S inner axes, for bmm1
862+ int64_t s2StrideSize = 0; // key Shape S inner axes, for bmm1
863+ int64_t maxS1Val = 0;
864+ int64_t maxS2Val = 0;
865+ 
866+ int64_t h1 = 0;
867+ int64_t h2 = 0;
868+ 
869+ int64_t s2sizeLimitMax = 1024;
870+ int64_t accumS1BlockNum = 0;
871+ 
872+ bool isSameAB = true;
873+ PromptFlashAttentionBaseApiTilingData baseApiTilingData;
874+ MLAGeneralTilingData mlaTilingData;
875+};
876+// end of class PromptFlashAttention
877+PFA_EXTERN_C ge::graphStatus TilingPromptFlashAttention(gert::TilingContext *context);
878+} // namespace optiling
879+ 
880+#endif // AIR_CXX_RUNTIME_V2_OP_IMPL_PROMPTFLASHATTENTION_H_
Acsrc/ops/prompt_flash_attention/op_host/prompt_flash_attention_tiling_compile_info.h+43-0
@@ -0,0 +1,43 @@
1+/**
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
3+ * MindIE is licensed under Mulan PSL v2.
4+ * You can use this software according to the terms and conditions of the Mulan PSL v2.
5+ * You may obtain a copy of Mulan PSL v2 at:
6+ * http://license.coscl.org.cn/MulanPSL2
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
8+ * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
9+ * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
10+ * See the Mulan PSL v2 for more details.
11+ */
12+ 
13+/*!
14+ * \file prompt_flash_attention_tiling_compile_info.h
15+ * \brief
16+ */
17+#ifndef PROMPT_FLASH_ATTENTION_TILING_STRUCT_COMPILE_INFO_H
18+#define PROMPT_FLASH_ATTENTION_TILING_STRUCT_COMPILE_INFO_H
19+#include <queue>
20+#include "exe_graph/runtime/tiling_context.h"
21+#include "op_host/data_copy_transpose_tiling_def.h"
22+#include "op_host/data_copy_transpose_tiling.h"
23+#include "register/tilingdata_base.h"
24+#include "tiling/tiling_api.h"
25+#include "register/op_def_registry.h"
26+ 
27+namespace optiling {
28+ 
29+struct PromptFlashAttentionCompileInfo {
30+ uint32_t aivNum;
31+ uint32_t aicNum;
32+ uint64_t ubSize;
33+ uint64_t l1Size;
34+ uint64_t l0CSize;
35+ uint64_t l0ASize;
36+ uint64_t l0BSize;
37+ size_t defaultSysWorkspaceSize;
38+ platform_ascendc::SocVersion socShortName;
39+};
40+ 
41+}
42+ 
43+#endif // PROMPT_FLASH_ATTENTION_TILING_STRUCT_COMPILE_INFO_H
Acsrc/ops/prompt_flash_attention/op_host/prompt_flash_attention_tiling_const.h+68-0
@@ -0,0 +1,68 @@
1+/**
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
3+ * MindIE is licensed under Mulan PSL v2.
4+ * You can use this software according to the terms and conditions of the Mulan PSL v2.
5+ * You may obtain a copy of Mulan PSL v2 at:
6+ * http://license.coscl.org.cn/MulanPSL2
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
8+ * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
9+ * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
10+ * See the Mulan PSL v2 for more details.
11+ */
12+ 
13+/*!
14+ * \file prompt_flash_attention_tiling_const.h
15+ * \brief
16+ */
17+#ifndef PROMPT_FLASH_ATTENTION_TILING_CONST_H
18+#define PROMPT_FLASH_ATTENTION_TILING_CONST_H
19+ 
20+namespace optiling {
21+constexpr uint32_t INT8SIZE = 1;
22+constexpr uint32_t UINT8SIZE = 1;
23+constexpr uint32_t FLOAT16SIZE = 2;
24+constexpr uint32_t FLOAT8SIZE = 1;
25+constexpr uint32_t BFLOAT16SIZE = 2;
26+constexpr uint32_t FLOAT32SIZE = 4;
27+constexpr uint32_t BOOLSIZE = 1;
28+ 
29+constexpr uint32_t NUM_16 = 16;
30+constexpr uint32_t NUM_32 = 32;
31+constexpr uint32_t FIRST_DIM = 0;
32+constexpr uint32_t SECOND_DIM = 1;
33+constexpr uint32_t THIRD_DIM = 2;
34+constexpr uint32_t FOURTH_DIM = 3;
35+constexpr uint32_t FIFTH_DIM = 4;
36+constexpr uint32_t DIM_0 = 0;
37+constexpr uint32_t DIM_1 = 1;
38+constexpr uint32_t DIM_2 = 2;
39+constexpr uint32_t DIM_3 = 3;
40+constexpr uint32_t DIM_4 = 4;
41+constexpr size_t DIM_NUM_1 = 1;
42+constexpr size_t DIM_NUM_3 = 3;
43+constexpr size_t DIM_NUM_4 = 4;
44+constexpr size_t DIM_NUM_5 = 5;
45+constexpr uint32_t N_SIZE_2 = 2;
46+constexpr uint32_t N_SIZE_4 = 4;
47+constexpr uint32_t N_SIZE_8 = 8;
48+constexpr uint32_t N_SIZE_16 = 16;
49+constexpr uint32_t N_SIZE_32 = 32;
50+constexpr uint32_t N_SIZE_64 = 64;
51+constexpr uint32_t N_SIZE_128 = 128;
52+constexpr uint32_t D_SIZE_192 = 192;
53+constexpr uint32_t D_SIZE_128 = 128;
54+constexpr uint32_t D_SIZE_64 = 64;
55+constexpr int HIGH_PRECISION = 0;
56+constexpr int HIGH_PERFORMANCE = 1;
57+constexpr int APPROXIMATE_COMPUTATION = 4;
58+constexpr uint32_t MSD_HIGH_PERFORMANCE_EXPEND_NUM = 2;
59+constexpr uint32_t MSD_HIGH_PRECISION_EXPEND_NUM = 3;
60+ 
61+constexpr uint32_t MAX_BATCH = 256U;
62+ 
63+constexpr int64_t MAX_VAR_LEN_SEQ_LEN = 4096L;
64+constexpr int64_t BALANCE_LOAD_LIST_SIZE = 8L;
65+constexpr int64_t COF[BALANCE_LOAD_LIST_SIZE] = {256, 384, 512, 640, 768, 896, 960, 1024};
66+} // namespace optiling
67+ 
68+#endif // PROMPT_FLASH_ATTENTION_TILING_CONST_H
Acsrc/ops/prompt_flash_attention/op_host/prompt_flash_attention_tiling_context.h+141-0
@@ -0,0 +1,141 @@
1+/**
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
3+ * MindIE is licensed under Mulan PSL v2.
4+ * You can use this software according to the terms and conditions of the Mulan PSL v2.
5+ * You may obtain a copy of Mulan PSL v2 at:
6+ * http://license.coscl.org.cn/MulanPSL2
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
8+ * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
9+ * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
10+ * See the Mulan PSL v2 for more details.
11+ */
12+ 
13+/*!
14+ * \file prompt_flash_attention_tiling_context.h
15+ * \brief
16+ */
17+#ifndef PROMPT_FLASH_ATTENTION_TILING_CONTEXT_H
18+#define PROMPT_FLASH_ATTENTION_TILING_CONTEXT_H
19+#include "op_host/data_copy_transpose_tiling_def.h"
20+#include "register/tilingdata_base.h"
21+#include "register/op_def_registry.h"
22+#include "prompt_flash_attention_tiling_compile_info.h"
23+ 
24+namespace optiling {
25+ 
26+/*
27+contextParams is a new structured defined for the use of FusedInferAttentionScore op.
28+It is meant to catch and organize all the necessary variables passed by FIAS tilling function.
29+It will be used as the input to the new 'runBigKernelWithParams' function in PFA tilling.
30+The old PFA tillingContext will also be transformed to this structure in the future.
31+*/
32+struct ContextParamsForPFATiling {
33+ const gert::Tensor *pseShift = nullptr;
34+ const gert::Tensor *attentionMask = nullptr;
35+ const gert::Tensor *actualSequenceLengthQ = nullptr;
36+ const gert::Tensor *actualSequenceLengthKV = nullptr;
37+ const gert::Tensor *antiquantScale = nullptr;
38+ const gert::Tensor *antiquantOffset = nullptr;
39+ const gert::Tensor *queryPaddingSize = nullptr;
40+ const gert::Tensor *kvPaddingSize = nullptr;
41+ const gert::Tensor *blockTable = nullptr;
42+ const gert::Tensor *keySharedPrefix = nullptr;
43+ const gert::Tensor *valueSharedPrefix = nullptr;
44+ const gert::Tensor *actualSharedPrefixLen = nullptr;
45+ const gert::Tensor *learnableSink = nullptr;
46+ 
47+ const gert::Tensor *keyAntiquantScale = nullptr;
48+ const gert::Tensor *valueAntiquantScale = nullptr;
49+ const gert::Tensor *KeyAntiquantOffset = nullptr;
50+ const gert::Tensor *valueAntiquantOffset = nullptr;
51+ const gert::Tensor *dequantScaleQuery = nullptr;
52+ 
53+ const gert::Tensor *qStartIdx = nullptr;
54+ const gert::Tensor *kvStartIdx = nullptr;
55+ 
56+ ge::DataType inputDataType = ge::DataType::DT_FLOAT16;
57+ ge::DataType kDataType = ge::DataType::DT_FLOAT16;
58+ ge::DataType vDataType = ge::DataType::DT_FLOAT16;
59+ ge::DataType qRopeDataType = ge::DataType::DT_FLOAT16;
60+ ge::DataType kRopeDataType = ge::DataType::DT_FLOAT16;
61+ ge::DataType pseShiftDataType = ge::DataType::DT_FLOAT16;
62+ ge::DataType maskDataType = ge::DataType::DT_FLOAT16;
63+ ge::DataType blockTableType = ge::DataType::DT_FLOAT16;
64+ ge::DataType outputDataType = ge::DataType::DT_FLOAT16;
65+ ge::DataType keySharedPrefixDataType = ge::DataType::DT_FLOAT16;
66+ ge::DataType valueSharedPrefixDataType = ge::DataType::DT_FLOAT16;
67+ ge::DataType learnableSinkDataType = ge::DataType::DT_FLOAT16;
68+ const char *opName = nullptr;
69+ const gert::StorageShape *queryInputShape = nullptr;
70+ const gert::StorageShape *keyInputShape = nullptr;
71+ const gert::StorageShape *queryRopeInputShape = nullptr;
72+ const gert::StorageShape *keyRopeInputShape = nullptr;
73+ const gert::StorageShape *valueInputShape = nullptr;
74+ const gert::StorageShape *pseShiftShape = nullptr;
75+ const gert::StorageShape *attentionMaskShape = nullptr;
76+ const gert::StorageShape *deqScale1Shape = nullptr;
77+ const gert::StorageShape *scale1Shape = nullptr;
78+ const gert::StorageShape *deqScale2Shape = nullptr;
79+ const gert::StorageShape *scale2Shape = nullptr;
80+ const gert::StorageShape *offset2Shape = nullptr;
81+ const gert::StorageShape *antiquantScaleShape = nullptr;
82+ const gert::StorageShape *antiquantOffsetShape = nullptr;
83+ const gert::StorageShape *blockTableShape = nullptr;
84+ const gert::StorageShape *outputShape = nullptr;
85+ const gert::StorageShape *lseoutputShape = nullptr;
86+ 
87+ const gert::StorageShape *dequantScaleQueryShape = nullptr;
88+ const gert::StorageShape *KeyAntiquantScaleShape = nullptr;
89+ const gert::StorageShape *valueAntiquantScaleShape = nullptr;
90+ const gert::StorageShape *KeyAntiquantOffsetShape = nullptr;
91+ const gert::StorageShape *valueAntiquantOffsetShape = nullptr;
92+ const gert::StorageShape *queryRope = nullptr;
93+ const gert::StorageShape *keyRope = nullptr;
94+ const gert::StorageShape *learnableSinkShape = nullptr;
95+ ge::DataType dequantScaleQueryType = ge::DataType::DT_FLOAT16;
96+ ge::DataType KeyAntiquantScaleType = ge::DataType::DT_FLOAT16;
97+ ge::DataType valueAntiquantScaleType = ge::DataType::DT_FLOAT16;
98+ ge::DataType KeyAntiquantOffsetType = ge::DataType::DT_FLOAT16;
99+ ge::DataType valueAntiquantOffsetType = ge::DataType::DT_FLOAT16;
100+ 
101+ const int64_t *innerPrecisePtr = nullptr;
102+ const int64_t *headsNumber = nullptr;
103+ const int32_t *sparseMode = nullptr;
104+ const int64_t *preToken = nullptr;
105+ const int64_t *nextToken = nullptr;
106+ const float *scaleValue = nullptr;
107+ const int32_t *blockSize = nullptr;
108+ const char *layout = nullptr;
109+ const int64_t *numKeyValueHeads = nullptr;
110+ size_t *workspaceSize = nullptr;
111+ const int64_t *pseType = nullptr;
112+ const PromptFlashAttentionCompileInfo *compileInfoPtr = nullptr;
113+ ge::DataType deqScaleType = ge::DataType::DT_FLOAT16;
114+ ge::DataType deqScale2Type = ge::DataType::DT_FLOAT16;
115+ ge::DataType quantScale2Type = ge::DataType::DT_FLOAT16;
116+ ge::DataType quantOffset2Type = ge::DataType::DT_FLOAT16;
117+ uint32_t isKvContinuous = 1;
118+ std::vector<const gert::StorageShape *> kTensorList = {nullptr};
119+ std::vector<const gert::StorageShape *> vTensorList = {nullptr};
120+ uint32_t maxKVs = 0;
121+ uint32_t fromFused = 0;
122+ uint32_t emptyTensor = 0;
123+ uint32_t isBSNDOut = 0;
124+ uint32_t transposeLayout = 0;
125+ const bool *softmaxLseFlag = nullptr;
126+ bool isSoftMaxLseEnable = false;
127+ uint32_t fromTilingSink =
128+ 0; // Flag indicating whether it is the step to enter the workspace calculation from tiling sinking
129+ bool hasKeyAntiquantScale = 0;
130+ bool hasValueAntiquantScale = 0;
131+ uint32_t isMsd = 0;
132+ const int64_t *keyAntiquantMode = nullptr;
133+ const int64_t *valueAntiquantMode = nullptr;
134+ const int64_t *queryQuantMode = nullptr;
135+ bool hasKeyAntiquantOffset = 0;
136+ bool hasLearnableSink = 0;
137+};
138+ 
139+} // namespace optiling
140+ 
141+#endif // PROMPT_FLASH_ATTENTION_TILING_CONTEXT_H
Acsrc/ops/prompt_flash_attention/op_host/prompt_flash_attention_tiling_struct.h+55-0
@@ -0,0 +1,55 @@
1+/**
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
3+ * MindIE is licensed under Mulan PSL v2.
4+ * You can use this software according to the terms and conditions of the Mulan PSL v2.
5+ * You may obtain a copy of Mulan PSL v2 at:
6+ * http://license.coscl.org.cn/MulanPSL2
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
8+ * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
9+ * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
10+ * See the Mulan PSL v2 for more details.
11+ */
12+ 
13+/*!
14+ * \file prompt_flash_attention_tiling_struct.h
15+ * \brief
16+ */
17+#ifndef PROMPT_FLASH_ATTENTION_TILING_STRUCT_H
18+#define PROMPT_FLASH_ATTENTION_TILING_STRUCT_H
19+namespace optiling {
20+ 
21+enum class InputLayout {
22+ SH,
23+ BSH,
24+ BNSD,
25+ NSD,
26+ BSND,
27+ BNSD_BSND,
28+ TND,
29+ NTD,
30+ NTD_TND,
31+ NZ,
32+ BBH,
33+ BNBD,
34+ NONE,
35+};
36+ 
37+enum class TilingMod {
38+ CVSAME = 0,
39+ CVDIFF,
40+ CVDIFF_BASE_API,
41+ CVDIFF_MLA,
42+};
43+ 
44+enum class SplitCoreMode {
45+ SPLIT_NBS_VECTOR = 0,
46+ SPLIT_NBS_CUBE,
47+ SPLIT_ONEN_VECTOR,
48+ SPLIT_ONEN_CUBE,
49+ BALANCE_VECTOR,
50+ BALANCE_CUBE,
51+ SPLIT_S1OUT_CUBE,
52+};
53+} // namespace optiling
54+ 
55+#endif // PROMPT_FLASH_ATTENTION_TILING_STRUCT_H
Acsrc/ops/prompt_flash_attention/op_host/prompt_flash_attention_tiling_v2.cpp+5518-0
Acsrc/ops/prompt_flash_attention/op_host/prompt_flash_attention_tiling_v2.h+414-0
@@ -0,0 +1,414 @@
1+/**
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
3+ * MindIE is licensed under Mulan PSL v2.
4+ * You can use this software according to the terms and conditions of the Mulan PSL v2.
5+ * You may obtain a copy of Mulan PSL v2 at:
6+ * http://license.coscl.org.cn/MulanPSL2
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
8+ * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
9+ * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
10+ * See the Mulan PSL v2 for more details.
11+ */
12+ 
13+/*!
14+ * \file prompt_flash_attention_tiling_v2.h
15+ * \brief
16+ */
17+#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_PROMPTFLASHATTENTION_V2_H_
18+#define AIR_CXX_RUNTIME_V2_OP_IMPL_PROMPTFLASHATTENTION_V2_H_
19+#include <unordered_map>
20+#include "op_host/data_copy_transpose_tiling_def.h"
21+#include "op_host/data_copy_transpose_tiling.h"
22+#include "register/tilingdata_base.h"
23+#include "tiling/tiling_api.h"
24+#include "register/op_def_registry.h"
25+#include "../../common/op_host/fia_tiling_templates_registry.h"
26+#include "prompt_flash_attention_tiling_context.h"
27+#include "prompt_flash_attention_tiling_struct.h"
28+#include "prompt_flash_attention_tiling_const.h"
29+#include "../../common/op_kernel/arch35/flash_attention_score_tiling_regbase.h"
30+#include "../op_kernel/arch35/prompt_flash_attention_tiling_regbase.h"
31+#include "../../common/op_host/fia_tiling_base.h"
32+ 
33+namespace optiling {
34+namespace v2 {
35+#define CHECK_D_LIMITED_SCENARIO(dSize) ((dSize) != 64 && (dSize) != 128)
36+ 
37+struct PFAShapeInfo {
38+ uint32_t b = 0;
39+ uint64_t n = 0;
40+ uint64_t s = 0;
41+ uint32_t d = 0;
42+ uint64_t h = 0;
43+ uint32_t t = 0;
44+};
45+static const std::unordered_map<ge::DataType, std::string> g_strDataTypePfa = {
46+ {ge::DT_FLOAT, "DT_FLOAT"},
47+ {ge::DT_FLOAT16, "DT_FLOAT16"},
48+ {ge::DT_INT8, "DT_INT8"},
49+ {ge::DT_INT16, "DT_INT16"},
50+ {ge::DT_UINT16, "DT_UINT16"},
51+ {ge::DT_UINT8, "DT_UINT8"},
52+ {ge::DT_INT32, "DT_INT32"},
53+ {ge::DT_INT64, "DT_INT64"},
54+ {ge::DT_UINT32, "DT_UINT32"},
55+ {ge::DT_UINT64, "DT_UINT64"},
56+ {ge::DT_BOOL, "DT_BOOL"},
57+ {ge::DT_DOUBLE, "DT_DOUBLE"},
58+ {ge::DT_STRING, "DT_STRING"},
59+ {ge::DT_DUAL_SUB_INT8, "DT_DUAL_SUB_INT8"},
60+ {ge::DT_DUAL_SUB_UINT8, "DT_DUAL_SUB_UINT8V"},
61+ {ge::DT_COMPLEX64, "DT_COMPLEX64"},
62+ {ge::DT_COMPLEX128, "DT_COMPLEX128"},
63+ {ge::DT_QINT8, "DT_QINT8"},
64+ {ge::DT_QINT16, "DT_QINT16"},
65+ {ge::DT_QINT32, "DT_QINT32"},
66+ {ge::DT_QUINT8, "DT_QUINT8"},
67+ {ge::DT_QUINT16, "DT_QUINT16"},
68+ {ge::DT_RESOURCE, "DT_RESOURCE"},
69+ {ge::DT_STRING_REF, "DT_STRING_REF"},
70+ {ge::DT_DUAL, "DT_DUAL"},
71+ {ge::DT_VARIANT, "DT_VARIANT"},
72+ {ge::DT_BF16, "DT_BF16"},
73+ {ge::DT_HIFLOAT8, "DT_HIFLOAT8"},
74+ {ge::DT_FLOAT8_E4M3FN, "DT_FLOAT8_E4M3FN"},
75+ {ge::DT_UNDEFINED, "DT_UNDEFINED"},
76+};
77+ 
78+static inline std::string GetPfaDataTypeStr(ge::DataType type) {
79+ ge::DataType findDype = (g_strDataTypePfa.find(type) == g_strDataTypePfa.end()) ? ge::DT_UNDEFINED : type;
80+ return g_strDataTypePfa.at(findDype);
81+}
82+ 
83+enum class AntiquantTypeEnum : uint8_t {
84+ PER_CHANNEL = 0,
85+ PER_TOKEN = 1,
86+ K_PER_CHANNEL_V_PER_TOKEN = 2,
87+ PER_TOKEN_HEAD = 3,
88+ PER_TOKEN_PAGE_ATTENTION = 4,
89+ PER_TOKEN_HEAD_PAGE_ATTENTION = 5,
90+ PER_BLOCK = 7
91+};
92+ 
93+class PromptFlashAttentionTilingV2 : public FiaTilingBase {
94+ public:
95+ explicit PromptFlashAttentionTilingV2(gert::TilingContext *context) : FiaTilingBase(context) {}
96+ ~PromptFlashAttentionTilingV2() override = default;
97+ ge::graphStatus RunBigKernelTilingWithParams(ContextParamsForPFATiling &contextKeyParams,
98+ uint32_t &numBlocksToBeSet, PromptFlashAttentionTilingDataV2 &tilingData);
99+ ge::graphStatus PromptFlashAttentionSetTilingData() const;
100+ bool CheckNonEmptyShapeExceptions(const ContextParamsForPFATiling &contextKeyParams,
101+ const gert::StorageShape *shape, const std::string &sName) const;
102+ ge::graphStatus DoSubOpTiling(
103+ PromptFlashAttentionTilingDataV2 &tilingData, ContextParamsForPFATiling &contextParamsForPFATiling);
104+ ge::graphStatus ConvertContextToPFAParams(ContextParamsForPFATiling &contextKeyParams) const;
105+ void SetTilingKey(ContextParamsForPFATiling &contextKeyParams) const;
106+ 
107+ protected:
108+ void InitializeMaxWorkspace(PFAShapeInfo &queryShapeInfo, PFAShapeInfo &keyShapeInfo,
109+ std::vector<int64_t> &actualSeqLengths, std::vector<int64_t> &actualSeqLengthsKV);
110+ void InitTilingInfo(TilingInfo *tilingInfo) override {}
111+ bool IsCapable() override { return true; }
112+ ge::graphStatus DoOpTiling() override;
113+ void PromptFlashAttentionInitOutputSplit(int64_t totalSize, PromptFlashAttentionTilingDataV2 &tilingData);
114+ ge::graphStatus CheckEmptyTensor(ContextParamsForPFATiling &contextKeyParams);
115+ void SetEmptyTensor(ContextParamsForPFATiling &contextKeyParams, uint32_t &numBlocksToBeSet,
116+ PromptFlashAttentionTilingDataV2 &tilingData);
117+ bool CheckIODataType(ContextParamsForPFATiling &contextKeyParams);
118+ bool SetInputLayout(const char *layout);
119+ bool GetAndCheckShape(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &shapeInfo,
120+ const gert::StorageShape *shape, const std::string &sName) const;
121+ bool GetAndCheckRopeShape(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &shapeInfo,
122+ PFAShapeInfo &ropeShapeInfo, const gert::StorageShape *shape, const std::string &sName,
123+ const std::string &rName) const;
124+ bool SetShape(ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *shape,
125+ const std::string inputName, int64_t &b, int64_t &n, int64_t &s, int64_t &d, int64_t &h, int64_t &t) const;
126+ bool CheckQueryOutParamsConsistency(const ContextParamsForPFATiling &contextKeyParams,
127+ const gert::StorageShape *queryShape, const gert::StorageShape *outShape) const;
128+ bool CheckKVDataType(ContextParamsForPFATiling &contextKeyParams) const;
129+ bool CheckRopeDataType(ContextParamsForPFATiling &contextKeyParams) const;
130+ bool CheckKeyValueParamsConsistency(ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *keyShape,
131+ const gert::StorageShape *valueShape);
132+ bool SetAndCheckHeadNumRatio(
133+ ContextParamsForPFATiling &contextKeyParams, PromptFlashAttentionTilingDataV2 &tilingData);
134+ bool CheckInputDimAndHeadNum(
135+ ContextParamsForPFATiling &contextKeyParams, const uint64_t nQAttr, const uint64_t nKVAttr);
136+ bool CheckPostQuantShape(const ContextParamsForPFATiling &contextKeyParams,
137+ const gert::StorageShape *quantScale2Shape, const gert::StorageShape *quantOffset2Shape,
138+ const PFAShapeInfo &queryShapeInfo, const PFAShapeInfo &valueShapeInfo) const;
139+ bool CheckPostQuantParams(const ContextParamsForPFATiling &contextKeyParams, const PFAShapeInfo &queryShapeInfo,
140+ const PFAShapeInfo &valueShapeInfo) const;
141+ bool CheckPerTensorQuantParams(
142+ const ContextParamsForPFATiling &contextKeyParams, const PFAShapeInfo &queryShapeInfo) const;
143+ bool CheckPerblockQuantParams(const ContextParamsForPFATiling &contextKeyParams, const PFAShapeInfo &queryShapeInfo,
144+ const PFAShapeInfo &keyShapeInfo, const PFAShapeInfo &valueShapeInfo) const;
145+ bool CheckAntiquantParamsShape(ContextParamsForPFATiling &contextKeyParams) const;
146+ bool GetAndCheckPrefixShape(
147+ ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &keyShapeInfo, PFAShapeInfo &prefixShapeInfo) const;
148+ bool CheckKeyValuePrefixConsistency(
149+ ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *keyShape);
150+ bool CheckActSharedPrefix(ContextParamsForPFATiling &contextKeyParams, const uint32_t sPrefix, const uint32_t sKV);
151+ bool CheckPAKeyValueShape(ContextParamsForPFATiling &contextKeyParams, int64_t &keyDim1,
152+ PFAShapeInfo &queryShapeInfo, const gert::StorageShape *keyShape, const gert::StorageShape *valueShape,
153+ const size_t keyDim, const int32_t *blockSize, int64_t blockNumValid, int64_t headNumRatio);
154+ bool CheckPACacheShape(ContextParamsForPFATiling &contextKeyParams, const size_t keyDim, PFAShapeInfo &shapeInfo,
155+ const gert::StorageShape *shape, const int32_t *blockSize, int64_t blockNumValid, int64_t headNumRatio,
156+ const std::string &sName);
157+ bool CheckBlockTableShape(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo,
158+ PFAShapeInfo &queryRopeShapeInfo, const int32_t *blockSize, const gert::StorageShape *blockTableShape,
159+ PromptFlashAttentionTilingDataV2 &tilingData);
160+ bool CheckMaskShape(ContextParamsForPFATiling &contextKeyParams, const int32_t *sparseMode, int64_t &attenMaskBatch,
161+ int64_t &attenMaskS1, int64_t &attenMaskS2, bool &checkMask, const uint64_t sQ, const uint64_t sK,
162+ const uint32_t batchSize, std::string &strMaskShape);
163+ void SetSparseModeData(ContextParamsForPFATiling &contextKeyParams, const gert::StorageShape *attenMaskShape,
164+ const int32_t *sparseMode, const int64_t *preTokens, const int64_t *nextTokens);
165+ bool CheckMaskShapeCrossSparse(ContextParamsForPFATiling &contextKeyParams,
166+ PromptFlashAttentionTilingDataV2 &tilingData, const int32_t *sparseMode, uint64_t sQ, const uint64_t sK,
167+ const uint32_t batchSize);
168+ bool CheckMaskCrossIFAMLA(ContextParamsForPFATiling &contextKeyParams, const int32_t *sparseMode, uint64_t queryS);
169+ bool CheckIO(
170+ ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo, PFAShapeInfo &valueShapeInfo);
171+ bool CheckKV(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &keyShapeInfo, PFAShapeInfo &valueShapeInfo);
172+ bool CheckQueryAndKey(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo,
173+ PFAShapeInfo &keyShapeInfo, PromptFlashAttentionTilingDataV2 &tilingData);
174+ bool CheckPFAMerge(ContextParamsForPFATiling &contextKeyParams, const PFAShapeInfo &queryShapeInfo) const;
175+ bool CheckRope(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo,
176+ PFAShapeInfo &keyShapeInfo, PFAShapeInfo &queryRopeShapeInfo);
177+ bool CheckLayout(ContextParamsForPFATiling &contextKeyParams) const;
178+ bool CheckIFAMLA(ContextParamsForPFATiling &contextKeyParams, const PFAShapeInfo &queryShapeInfo) const;
179+ bool CheckQuant(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo,
180+ PFAShapeInfo &keyShapeInfo, const PFAShapeInfo &valueShapeInfo) const;
181+ bool CheckQScaleShape4MLAFullQuant(ContextParamsForPFATiling &contextKeyParams) const;
182+ bool CheckKVScaleShape4MLAFullQuant(ContextParamsForPFATiling &contextKeyParams) const;
183+ bool CheckMLAFullQuant(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo) const;
184+ bool CheckPrefix(ContextParamsForPFATiling &contextKeyParams, const PFAShapeInfo &queryShapeInfo,
185+ PFAShapeInfo &keyShapeInfo, PromptFlashAttentionTilingDataV2 &tilingData);
186+ bool CheckActSeq(const ContextParamsForPFATiling &contextKeyParams, const PFAShapeInfo &queryShapeInfo) const;
187+ bool CheckActSeqLen(
188+ ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo, PFAShapeInfo &keyShapeInfo);
189+ bool CheckPATypeAndShape(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo,
190+ PFAShapeInfo &queryRopeShapeInfo, PromptFlashAttentionTilingDataV2 &tilingData);
191+ bool CheckPseShiftTypeAndShape(
192+ ContextParamsForPFATiling &contextKeyParams, uint32_t b, uint32_t n, uint32_t s1, uint32_t s2);
193+ bool CheckInnerPrecise(ContextParamsForPFATiling &contextKeyParams, PromptFlashAttentionTilingDataV2 &tilingData);
194+ bool CheckMaskTypeAndShape(
195+ ContextParamsForPFATiling &contextKeyParams, PromptFlashAttentionTilingDataV2 &tilingData);
196+ void SetSparseType(uint32_t qS);
197+ bool CheckSparseMode(ContextParamsForPFATiling &contextKeyParams, uint32_t qS);
198+ bool CheckPACrossover(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo);
199+ bool CheckMaskCrossover(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo,
200+ PromptFlashAttentionTilingDataV2 &tilingData);
201+ bool CheckTNDLayoutCrossover(ContextParamsForPFATiling &contextKeyParams);
202+ bool CheckNTDLayoutCrossover(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo);
203+ bool CheckTransposeLayoutCrossover(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo);
204+ bool CheckLearnSink(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo);
205+ bool ParseActualSeqLengths(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo,
206+ std::vector<int64_t> &actualSeqLengths, std::vector<int64_t> &actualSeqLengthsKV);
207+ bool CheckMultiFeatureCrossover(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo,
208+ std::vector<int64_t> &actualSeqLengths, std::vector<int64_t> &actualSeqLengthsKV,
209+ PromptFlashAttentionTilingDataV2 &tilingData);
210+ bool CheckPerblockCrossover(
211+ ContextParamsForPFATiling &contextKeyParams, PromptFlashAttentionTilingDataV2 &tilingData);
212+ void SetTilingDataAttribute(
213+ ContextParamsForPFATiling &contextKeyParams, PromptFlashAttentionTilingDataV2 &tilingData);
214+ void GetEnableDN(ContextParamsForPFATiling &contextKeyParams, PromptFlashAttentionTilingDataV2 &tilingData,
215+ PFAShapeInfo &queryShapeInfo, PFAShapeInfo &valueShapeInfo, std::vector<int64_t> &actualSeqLengths,
216+ std::vector<int64_t> &actualSeqLengthsKV);
217+ void SetTilingData(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo,
218+ PFAShapeInfo &queryRopeShapeInfo, PFAShapeInfo &valueShapeInfo, PromptFlashAttentionTilingDataV2 &tilingData);
219+ void InferTilingMod(const ContextParamsForPFATiling &contextKeyParams, std::vector<int64_t> &actualSeqLengths,
220+ std::vector<int64_t> &actualSeqLengthsKV, uint32_t actualSeqArrayLen, uint32_t d);
221+ int64_t GetActualInnerBlockNums(int64_t sInnerIndexStart, int64_t sInnerIndexEnd, int64_t innerBlockNums) const;
222+ int64_t SumOfArithmeticSeries(int64_t an, int64_t d) const;
223+ int64_t GetCutBlockNums(
224+ int64_t blockSeqLengthKV, int64_t blockSeqLength, int64_t sInner, int64_t sOuter, int64_t token) const;
225+ void FixParamWithRowInvalid(
226+ int64_t &actualSeqLength, int64_t actualSeqLengthKV, int64_t &preTokensLeftUp, int64_t &nextTokensLeftUp) const;
227+ int64_t GetCalcBlockNumsOneHead(int64_t actualSeqLength, int64_t actualSeqLengthKV, uint32_t sOuterSize,
228+ uint32_t sInnerSize, int64_t preTokensLeftUp, int64_t nextTokensLeftUp, bool isAttenMaskUsed) const;
229+ void ComputeSplitNBSeq(PromptFlashAttentionTilingDataV2 &tilingData, uint32_t batchSize,
230+ const size_t tilingElementArrayLen, std::vector<int64_t> &actualSeqLengths,
231+ std::vector<int64_t> &actualSeqLengthsKV, uint32_t sOuterSize, uint32_t sInnerSize, double coreWightTarget,
232+ uint32_t &curCore);
233+ void PromptFlashAttentionSplitNBSeq(PromptFlashAttentionTilingDataV2 &tilingData,
234+ std::vector<int64_t> &actualSeqLengths, std::vector<int64_t> &actualSeqLengthsKV, bool isAttenMaskUsed);
235+ void InferSplitCoreMode();
236+ void InferConstantization();
237+ bool AdjustCVTilingCVDiff(const ContextParamsForPFATiling &contextKeyParams, uint32_t &sOuterFactor,
238+ uint32_t &sInnerFactor, uint32_t &softmaxSOuterFactor, PromptFlashAttentionTilingDataV2 &tilingData,
239+ const PFAShapeInfo &queryShapeInfo);
240+ void GetMatMulType(matmul_tiling::DataType &mmInputType, matmul_tiling::DataType &mmOutputType);
241+ bool EnableMTE2BmmPipe(PromptFlashAttentionTilingDataV2 &tilingData, matmul_tiling::MatmulApiTiling &bmm,
242+ TCubeTiling &bmmTilingData, uint32_t sOuterFactor, uint32_t sInnerFactor);
243+ void EnableBmmDoubleBuffer(TCubeTiling &bmmTilingData) const;
244+ bool PromptFlashAttentionCheckBmm1(PromptFlashAttentionTilingDataV2 &tilingData, TCubeTiling &bmm1TilingData,
245+ int64_t l1SizeRemain, int64_t l0CSize, uint32_t sOuterFactor, uint32_t sInnerFactor, bool autoBaseMNK = false);
246+ bool PromptFlashAttentionCheckBmm2(PromptFlashAttentionTilingDataV2 &tilingData, TCubeTiling &bmm2TilingData,
247+ int64_t l1SizeRemain, int64_t l0CSize, uint32_t sOuterFactor, uint32_t sInnerFactor, uint32_t dSplitFactor,
248+ bool autoBaseMNK = false);
249+ bool PromptFlashAttentionComputeCVDiffParams(PromptFlashAttentionTilingDataV2 &tilingData, int64_t l1Size,
250+ int64_t l0CSize, uint32_t &sOuterFactor, uint32_t &sInnerFactor);
251+ void GetPreNextTokensLeftUp(PromptFlashAttentionTilingDataV2 &tilingData, int64_t actualSeqLength,
252+ int64_t actualSeqLengthKV, int64_t &preTokensLeftUp, int64_t &nextTokensLeftUp);
253+ void PromptFlashAttentionInitSoftmaxLseOutputSplit(int64_t totalSize, PromptFlashAttentionTilingDataV2 &tilingData);
254+ bool TilingGetTilingKeyAttentionAscendC(
255+ ContextParamsForPFATiling &contextKeyParams, PromptFlashAttentionTilingDataV2 &tilingData);
256+ size_t GetPFAWorkSpaceSize(PromptFlashAttentionTilingDataV2 &tilingData);
257+ ge::graphStatus SetPlatMemoryInfo(ContextParamsForPFATiling &contextKeyParams);
258+ ge::graphStatus SetAttributeInfo(ContextParamsForPFATiling &contextKeyParams);
259+ ge::graphStatus CheckTensorInvalid(const ContextParamsForPFATiling &contextKeyParams) const;
260+ ge::graphStatus CheckSingleAttribute(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo,
261+ PFAShapeInfo &keyShapeInfo, PFAShapeInfo &valueShapeInfo, PFAShapeInfo &queryRopeShapeInfo,
262+ PromptFlashAttentionTilingDataV2 &tilingData);
263+ ge::graphStatus CheckCrossoverAttribute(ContextParamsForPFATiling &contextKeyParams, PFAShapeInfo &queryShapeInfo,
264+ PromptFlashAttentionTilingDataV2 &tilingData);
265+ ge::graphStatus AdjustTilingData(ContextParamsForPFATiling &contextKeyParams,
266+ PromptFlashAttentionTilingDataV2 &tilingData, const PFAShapeInfo &queryShapeInfo);
267+ ge::graphStatus ComputeTilingData(ContextParamsForPFATiling &contextKeyParams,
268+ std::vector<int64_t> &actualSeqLengths, std::vector<int64_t> &actualSeqLengthsKV,
269+ PromptFlashAttentionTilingDataV2 &tilingData);
270+ ge::graphStatus ComputeTilingKey(ContextParamsForPFATiling &contextKeyParams, uint32_t &numBlocksToBeSet,
271+ PromptFlashAttentionTilingDataV2 &tilingData);
272+ void SetAttenMaskCompressMode();
273+ void SetLayoutType();
274+ void PFATilingDataconvert(PromptFlashAttentionTilingDataV2 &tilingData);
275+ void SetMultiCoreParamsRegbase(int64_t totalSize, int64_t coreNum);
276+ bool IsFlashDecode(ContextParamsForPFATiling &contextKeyParams, uint64_t bng) const;
277+ ge::graphStatus SplitBNS(PromptFlashAttentionTilingDataV2 &tilingData, uint64_t bng);
278+ bool CheckAlibiPseShiftTypeAndShape(ContextParamsForPFATiling &contextKeyParams, uint32_t n);
279+ ge::graphStatus SetQKVStartIdx(ContextParamsForPFATiling &contextKeyParams);
280+ bool CheckAlibiPseCrossover(ContextParamsForPFATiling &contextKeyParams);
281+ void GetMaxWorkspaceFlag(ContextParamsForPFATiling &contextKeyParams);
282+ void GetQueryDimAndOutDim(const gert::StorageShape *queryShape, const gert::StorageShape *outShape,
283+ const std::string &layoutStr, int64_t &tmpqueryDim, int64_t &outDim, uint32_t i) const;
284+ 
285+ void UpdateTilingKeyLayoutType();
286+ void UpdateTilingKeyConfig(
287+ ContextParamsForPFATiling &contextKeyParams, PromptFlashAttentionTilingDataV2 &tilingData);
288+ void UpdateTilingKeyPseMode();
289+ void UpdateTilingKeyQuantMode(ge::DataType inputDataType);
290+ void UpdateTilingKeyAttenMask(ge::DataType inputDataType);
291+ void UpdateTilingKeyHasRope(ge::DataType inputDataType);
292+ void UpdateTilingKeyIsPa(ge::DataType inputDataType);
293+ void UpdateTilingKeyIsFd(ge::DataType inputDataType);
294+ void UpdateTilingKeyEmptyTensor();
295+ void UpdateTilingKeyEnableKVPrefix();
296+ void UpdateTilingKeySplitCoreMode();
297+ 
298+ public:
299+ uint8_t inOutLayoutType = 0;
300+ uint16_t config = 0;
301+ uint8_t pseMode = 0;
302+ uint8_t quantMode = 0;
303+ bool hasAttenMask = false;
304+ bool hasRope = false;
305+ bool isPa = false;
306+ bool isFd = false;
307+ bool emptyTensor = false;
308+ bool enableKVPrefix = false;
309+ bool enableS1OutSplit = false;
310+ 
311+ protected:
312+ ContextParamsForPFATiling *contextKeyParamsPtr = nullptr;
313+ int64_t ubSizeRemain = 1;
314+ bool enableFlashDecode = false;
315+ bool isSOuterNoTail = true;
316+ bool isSInnerNoTail = true;
317+ bool isDNoTail = true;
318+ bool enableTensorList = false;
319+ bool enableLeftPadding = false;
320+ bool enableActSeqLen = false;
321+ bool enableActSeqLenKV = false;
322+ bool enableKVAntiquant = false;
323+ bool enablePseShift = false;
324+ bool enableAlibiPse = false;
325+ bool enableMask = false;
326+ bool enableQuantBF16 = false;
327+ bool enableMatmulNorm = false;
328+ bool enablePA = false;
329+ bool enableSplitSeqOneN = false;
330+ bool isDefaultSparseMode = false;
331+ bool isKVHasPrefix = false;
332+ bool isBandMode = false;
333+ bool enableIFAMLA = false;
334+ bool enableIFA = false;
335+ bool enableIFAMask = false;
336+ bool enableIFAMLAFullQuant = false;
337+ // MLPerf合轴优化
338+ bool enablePFAMerge = false;
339+ uint32_t pfaMergeGLimit = 16;
340+ uint32_t pfaMergeQsLimit = 4;
341+ bool enablePFAMLA = false;
342+ bool enablePFARope = false;
343+ bool enableDN = false;
344+ bool enablePostQuant = false;
345+ bool enablePertensorQuant = false;
346+ bool enablePerblockQuant = false;
347+ // attention sink
348+ bool enableLearnSink = false;
349+ uint64_t gSize = 1;
350+ int64_t t1Size = 0;
351+ int64_t t2Size = 0;
352+ InputLayout inputLayout = InputLayout::BSH;
353+ ge::DataType inputType{ge::DT_FLOAT16};
354+ ge::DataType outputType{ge::DT_FLOAT16};
355+ ge::DataType pseShiftElemType{ge::DT_FLOAT16};
356+ ge::DataType queryType{ge::DT_FLOAT};
357+ ge::DataType keyType{ge::DT_FLOAT};
358+ ge::DataType valueType{ge::DT_FLOAT};
359+ uint32_t dataTypeSize = FLOAT32SIZE;
360+ uint32_t outputDataTypeSize = FLOAT32SIZE;
361+ uint32_t maskElemSize = FLOAT32SIZE;
362+ int32_t ifaBlockSizeBase = 32;
363+ uint32_t coreNum = 0;
364+ uint32_t aivNum = 0;
365+ uint32_t aicNum = 0;
366+ uint32_t typeByteNum = 0;
367+ uint32_t outputTypeByteNum = 0;
368+ uint32_t softmaxTypeByteNum = 0;
369+ uint32_t pseShiftTypeByteNum = 0;
370+ uint32_t pseShiftElemSize = 0;
371+ uint32_t pseMaskMaxSize = 0;
372+ int64_t pseShiftBatch = 0;
373+ int64_t pseShiftS1 = 0;
374+ int64_t pseShiftS2 = 0;
375+ int64_t actSeqLenDims = 0;
376+ int64_t actSeqLenKVDims = 0;
377+ int64_t middleActualSeqLengths = 0;
378+ int64_t actualSharedPrefixLen = 0;
379+ uint32_t needInit = 0U;
380+ uint32_t usePseShift = 0;
381+ // There is no S2 axis for PA. Use the change amount to normalize the S2 length in both PA and non PA scenarios
382+ uint32_t S2 = 0;
383+ int32_t blockTableDim2 = 1;
384+ int32_t paBlockNumSum = 1;
385+ uint32_t maskTypeByteNum = 0;
386+ uint32_t softmaxDataTypeSize = FLOAT32SIZE; // BF16 calculates through FP32
387+ platform_ascendc::SocVersion curShortSocName;
388+ uint32_t layoutType = 0;
389+ uint32_t paLayoutType = 0;
390+ int64_t sparsePreTokens = 0;
391+ int64_t sparseNextTokens = 0;
392+ int32_t sparseModeVal = 0;
393+ int64_t maxActualseqKV = 0;
394+ SplitCoreMode splitCoreMode = SplitCoreMode::SPLIT_NBS_VECTOR;
395+ bool isConstantization = false;
396+ uint32_t splitS2 = 1; // It can only be 0 when the D axis is split
397+ int32_t innerPrecise = HIGH_PERFORMANCE;
398+ uint32_t sOuterFactorTiling = 0;
399+ uint32_t softmaxSInnerFactorTiling = 0;
400+ uint32_t softmaxSOuterFactorTiling = 0;
401+ matmul_tiling::PlatformInfo ascendPlatformInfo;
402+ 
403+ bool isMaxWorkspace = false;
404+ bool isQKVDDifferent = false;
405+ bool faRunFlag_ = true;
406+ uint8_t attenMaskShapeType = 0; // 0: (B,N2,G,S1,S2), 1: (B,1,1,S1,S2), 2: (1,1,1,S1,S2)
407+ uint8_t sparseType = 0;
408+ int64_t pseType = 0;
409+ FlashAttentionScoreSimplifiedTilingData faTilingAdapter;
410+};
411+} // namespace v2
412+} // namespace optiling
413+ 
414+#endif // AIR_CXX_RUNTIME_V2_OP_IMPL_PROMPTFLASHATTENTION_V2_H_