Star1130
2739
代码
代码
Issues217
Pull Requests406
Actions
项目讨论
Wiki
分析
Star1130
2739
已合并
update FA/FAG #856
huanyongben创建于 1月20日
update FA/FAG #856
已合并
huanyongben创建于 1月20日
从huanyongben/ops-transformermaster合入到cann/ops-transformermaster
共 96 个文件变更+12506-1049
Mattention/flash_attention_score/README.md+8-6
@@ -22,13 +22,11 @@
22 注意力的正向计算公式如下：22 注意力的正向计算公式如下：
23 23 
24 - pseType=1时，公式如下：24 - pseType=1时，公式如下：
25- 
26 $$25 $$
27 attention\_out=Dropout(Softmax(Mask(scale*(pse+(query*d\_scale\_q)*(key*d\_scale\_k)^T), atten\_mask)), keep\_prob)*(value*d\_scale\_v)26 attention\_out=Dropout(Softmax(Mask(scale*(pse+(query*d\_scale\_q)*(key*d\_scale\_k)^T), atten\_mask)), keep\_prob)*(value*d\_scale\_v)
28 $$27 $$
29 28 
30 - pseType≠1时，公式如下：29 - pseType≠1时，公式如下：
31- 
32 $$30 $$
33 attention\_out=Dropout(Softmax(Mask(scale*((query*d\_scale\_q)*(key*d\_scale\_k)^T) + pse),atten\_mask),keep\_prob)*(value*d\_scale\_v)31 attention\_out=Dropout(Softmax(Mask(scale*((query*d\_scale\_q)*(key*d\_scale\_k)^T) + pse),atten\_mask),keep\_prob)*(value*d\_scale\_v)
34 $$32 $$
@@ -55,21 +53,21 @@
55 <td>query</td>53 <td>query</td>
56 <td>输入</td>54 <td>输入</td>
57 <td>公式中的输入query。</td>55 <td>公式中的输入query。</td>
58- <td>BFLOAT16、FLOAT16、FLOAT</td>56+ <td>BFLOAT16、FLOAT16、FLOAT、FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8</td>
59 <td>ND</td>57 <td>ND</td>
60 </tr>58 </tr>
61 <tr>59 <tr>
62 <td>key</td>60 <td>key</td>
63 <td>输入</td>61 <td>输入</td>
64 <td>公式中的输入key。</td>62 <td>公式中的输入key。</td>
65- <td>BFLOAT16、FLOAT16、FLOAT</td>63+ <td>BFLOAT16、FLOAT16、FLOAT、FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8</td>
66 <td>ND</td>64 <td>ND</td>
67 </tr>65 </tr>
68 <tr>66 <tr>
69 <td>value</td>67 <td>value</td>
70 <td>输入</td>68 <td>输入</td>
71 <td>公式中的输入value。</td>69 <td>公式中的输入value。</td>
72- <td>BFLOAT16、FLOAT16、FLOAT</td>70+ <td>BFLOAT16、FLOAT16、FLOAT、FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8</td>
73 <td>ND</td>71 <td>ND</td>
74 </tr>72 </tr>
75 <tr>73 <tr>
@@ -174,6 +172,9 @@
174 </tbody>172 </tbody>
175</table>173</table>
176 174 
175+- <term>Atlas A2 训练系列产品</term>：不支持FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8三种数据类型。
176+- <term>Atlas A3 训练系列产品</term>：不支持FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8三种数据类型。
177+ 
177## 约束说明178## 约束说明
178 179 
179- 输入query、key、value、realShiftOptional的数据类型必须一致。180- 输入query、key、value、realShiftOptional的数据类型必须一致。
@@ -182,7 +183,8 @@
182 - B：取值范围为1\~2M。当prefixOptional的时候B最大支持2K。183 - B：取值范围为1\~2M。当prefixOptional的时候B最大支持2K。
183 - N：取值范围为1\~256。184 - N：取值范围为1\~256。
184 - S：取值范围为1\~1M。185 - S：取值范围为1\~1M。
185- - D：取值范围为1\~512。186+ - D：取值范围为1\~768。输入query、key、value类型为FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8时，D取值范围为1\~128。
187+- 输入query、key、value类型为FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8时, 不支持queryRopeOptional、keyRopeOptional、realShiftOptional、attenMaskOptional、dropMaskOptional、keepProb、pseType等相关可选参数。
186- keepProb的取值范围为(0, 1]。188- keepProb的取值范围为(0, 1]。
187- 部分场景下，如果计算量过大可能会导致算子执行超时(aicore error类型报错，errorStr为：timeout or trap error)，此时建议做轴切分处理，注：这里的计算量会受B、S、N、D等参数的影响，值越大计算量越大。189- 部分场景下，如果计算量过大可能会导致算子执行超时(aicore error类型报错，errorStr为：timeout or trap error)，此时建议做轴切分处理，注：这里的计算量会受B、S、N、D等参数的影响，值越大计算量越大。
188- pseType为2或3的时候，当前只支持Sq和Skv等长。190- pseType为2或3的时候，当前只支持Sq和Skv等长。
Aattention/flash_attention_score/docs/aclnnFlashAttentionScoreV4.md+846-0文件内容审核中，请稍后刷新重试
Mattention/flash_attention_score/op_api/aclnn_flash_attention_score.cpp+240-64
@@ -98,6 +98,14 @@ struct FaShapeInfo {
98 bool needPadValue = false;98 bool needPadValue = false;
99};99};
100 100 
101+static bool StrideLimited() {
102+ NpuArch npuArch = GetCurrentPlatformInfo().GetCurNpuArch();
103+ if (npuArch == NpuArch::DAV_2201) {
104+ return true;
105+ }
106+ return false;
107+}
108+ 
101static void AnalysisAxisForBsh(const Shape &qShape, const Shape &kShape, const Shape &vShape, FaShapeInfo &shapeInfo)109static void AnalysisAxisForBsh(const Shape &qShape, const Shape &kShape, const Shape &vShape, FaShapeInfo &shapeInfo)
102{110{
103 shapeInfo.inputLayout = InputLayout::BSH;111 shapeInfo.inputLayout = InputLayout::BSH;
@@ -388,18 +396,22 @@ static bool isSupportMLA(const FaShapeInfo &shapeInfo, const aclTensor *query)
388}396}
389 397 
390static aclnnStatus InputDtypeCheck(const aclTensor *query, const aclTensor *key, const aclTensor *value,398static aclnnStatus InputDtypeCheck(const aclTensor *query, const aclTensor *key, const aclTensor *value,
391- const aclTensor *realShiftOptional, int64_t pseType, const aclTensor *sinkOptional)399+ const aclTensor *attentionOut, const aclTensor *realShiftOptional,
400+ int64_t pseType, const aclTensor *sinkOptional)
392{401{
393 auto vDtype = value->GetDataType();402 auto vDtype = value->GetDataType();
394 auto kDtype = key->GetDataType();403 auto kDtype = key->GetDataType();
395 auto qDtype = query->GetDataType();404 auto qDtype = query->GetDataType();
405+ auto outDtype = attentionOut->GetDataType();
396 if (qDtype != kDtype || kDtype != vDtype) {406 if (qDtype != kDtype || kDtype != vDtype) {
397 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The data type of query[%s], key[%s], value[%s] are not equal.",407 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The data type of query[%s], key[%s], value[%s] are not equal.",
398 op::ToString(DataType(qDtype)).GetString(), op::ToString(DataType(kDtype)).GetString(),408 op::ToString(DataType(qDtype)).GetString(), op::ToString(DataType(kDtype)).GetString(),
399 op::ToString(DataType(vDtype)).GetString());409 op::ToString(DataType(vDtype)).GetString());
400 return ACLNN_ERR_PARAM_INVALID;410 return ACLNN_ERR_PARAM_INVALID;
401 }411 }
402- if (!(qDtype == op::DataType::DT_FLOAT || qDtype == op::DataType::DT_FLOAT16 || qDtype == op::DataType::DT_BF16)) {412+ 
413+ if (StrideLimited() &&
414+ !(qDtype == op::DataType::DT_FLOAT || qDtype == op::DataType::DT_FLOAT16 || qDtype == op::DataType::DT_BF16)) {
403 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The data type of query/key/value is [%s], should be fp16, bf16 or fp32.",415 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The data type of query/key/value is [%s], should be fp16, bf16 or fp32.",
404 op::ToString(DataType(qDtype)).GetString());416 op::ToString(DataType(qDtype)).GetString());
405 return ACLNN_ERR_PARAM_INVALID;417 return ACLNN_ERR_PARAM_INVALID;
@@ -421,10 +433,10 @@ static aclnnStatus InputDtypeCheck(const aclTensor *query, const aclTensor *key,
421 }433 }
422 if (realShiftOptional != nullptr) {434 if (realShiftOptional != nullptr) {
423 auto pseDtype = realShiftOptional->GetDataType();435 auto pseDtype = realShiftOptional->GetDataType();
424- if (pseDtype != qDtype) {436+ if (pseDtype != outDtype) {
425 OP_LOGE(ACLNN_ERR_PARAM_INVALID,437 OP_LOGE(ACLNN_ERR_PARAM_INVALID,
426- "The data type %s of pse is not equal to the data type %s of query, key and value.",438+ "The data type %s of pse is not equal to the data type %s of attentionOut.",
427- op::ToString(DataType(pseDtype)).GetString(), op::ToString(DataType(qDtype)).GetString());439+ op::ToString(DataType(pseDtype)).GetString(), op::ToString(DataType(outDtype)).GetString());
428 return ACLNN_ERR_PARAM_INVALID;440 return ACLNN_ERR_PARAM_INVALID;
429 }441 }
430 }442 }
@@ -454,7 +466,7 @@ static inline bool CheckFormat(
454 if (!formatValid) {466 if (!formatValid) {
455 OP_LOGE(467 OP_LOGE(
456 ACLNN_ERR_PARAM_INVALID,468 ACLNN_ERR_PARAM_INVALID,
457- "Input and output format only support [ND]. Actual: query:[%s], key:[%s], value:[%s], softmaxMaxOut:[%s], softmaxSumOut:[%s], attentionOutOut:[%s].",469+ "Input and output format do not support [NZ]. Actual: query:[%s], key:[%s], value:[%s], softmaxMaxOut:[%s], softmaxSumOut:[%s], attentionOutOut:[%s].",
458 op::ToString(query->GetStorageFormat()).GetString(), op::ToString(key->GetStorageFormat()).GetString(),470 op::ToString(query->GetStorageFormat()).GetString(), op::ToString(key->GetStorageFormat()).GetString(),
459 op::ToString(value->GetStorageFormat()).GetString(), op::ToString(softmaxMaxOut->GetStorageFormat()).GetString(), 471 op::ToString(value->GetStorageFormat()).GetString(), op::ToString(softmaxMaxOut->GetStorageFormat()).GetString(),
460 op::ToString(softmaxSumOut->GetStorageFormat()).GetString(), op::ToString(attentionOutOut->GetStorageFormat()).GetString());472 op::ToString(softmaxSumOut->GetStorageFormat()).GetString(), op::ToString(attentionOutOut->GetStorageFormat()).GetString());
@@ -482,7 +494,7 @@ static inline bool CheckFormat(
482 formatValid = (formatValid && sinkOptional->GetStorageFormat() != op::Format::FORMAT_FRACTAL_NZ);494 formatValid = (formatValid && sinkOptional->GetStorageFormat() != op::Format::FORMAT_FRACTAL_NZ);
483 }495 }
484 if (!formatValid) {496 if (!formatValid) {
485- OP_LOGE(ACLNN_ERR_PARAM_INVALID,"Optional input format only support [ND]. Actual: queryRope:[%s], keyRope:[%s], realShiftOptional:[%s], "497+ OP_LOGE(ACLNN_ERR_PARAM_INVALID,"Optional input format do not support [NZ]. Actual: queryRope:[%s], keyRope:[%s], realShiftOptional:[%s], "
486 "dropMaskOptional:[%s], paddingMaskOptional:[%s], attenMaskOptional:[%s], sinkOptional:[%s].",op::ToString(queryRope->GetStorageFormat()).GetString(),498 "dropMaskOptional:[%s], paddingMaskOptional:[%s], attenMaskOptional:[%s], sinkOptional:[%s].",op::ToString(queryRope->GetStorageFormat()).GetString(),
487 op::ToString(keyRope->GetStorageFormat()).GetString(),op::ToString(realShiftOptional->GetStorageFormat()).GetString(),499 op::ToString(keyRope->GetStorageFormat()).GetString(),op::ToString(realShiftOptional->GetStorageFormat()).GetString(),
488 op::ToString(dropMaskOptional->GetStorageFormat()).GetString(),op::ToString(paddingMaskOptional->GetStorageFormat()).GetString(),500 op::ToString(dropMaskOptional->GetStorageFormat()).GetString(),op::ToString(paddingMaskOptional->GetStorageFormat()).GetString(),
@@ -525,6 +537,10 @@ static aclnnStatus AnalysisInput(const aclTensor *query, const aclTensor *key, c
525 return ACLNN_SUCCESS;537 return ACLNN_SUCCESS;
526 }538 }
527 539 
540+ if (!StrideLimited()) {
541+ return ACLNN_SUCCESS;
542+ }
543+ 
528 int64_t alignDim = (shapeInfo.axes.d < PAD_LOWER_BOUND_196 || shapeInfo.axes.d == PAD_ALIGN_SPL_SHAPE) ?544 int64_t alignDim = (shapeInfo.axes.d < PAD_LOWER_BOUND_196 || shapeInfo.axes.d == PAD_ALIGN_SPL_SHAPE) ?
529 PAD_BASIC_BLOCK : PAD_ALIGN_128;545 PAD_BASIC_BLOCK : PAD_ALIGN_128;
530 if (shapeInfo.axes.d % alignDim != 0 || shapeInfo.axes.dv % alignDim != 0) {546 if (shapeInfo.axes.d % alignDim != 0 || shapeInfo.axes.dv % alignDim != 0) {
@@ -603,7 +619,8 @@ static aclnnStatus Contiguous(const aclTensor *&query, const aclTensor *&key, co
603 const aclTensor *&realShiftOptional, const aclTensor *&dropMaskOptional,619 const aclTensor *&realShiftOptional, const aclTensor *&dropMaskOptional,
604 const aclTensor *&paddingMaskOptional, const aclTensor *&attenMaskOptional,620 const aclTensor *&paddingMaskOptional, const aclTensor *&attenMaskOptional,
605 const aclTensor *&queryRope, const aclTensor *&keyRope, const aclTensor *&sinkOptional,621 const aclTensor *&queryRope, const aclTensor *&keyRope, const aclTensor *&sinkOptional,
606- aclOpExecutor *executor)622+ const aclTensor *&dScaleQOptional, const aclTensor *&dScaleKOptional,
623+ const aclTensor *&dScaleVOptional, aclOpExecutor *executor)
607{624{
608 query = l0op::Contiguous(query, executor);625 query = l0op::Contiguous(query, executor);
609 CHECK_RET(query != nullptr, ACLNN_ERR_PARAM_NULLPTR);626 CHECK_RET(query != nullptr, ACLNN_ERR_PARAM_NULLPTR);
@@ -639,12 +656,27 @@ static aclnnStatus Contiguous(const aclTensor *&query, const aclTensor *&key, co
639 sinkOptional = l0op::Contiguous(sinkOptional, executor);656 sinkOptional = l0op::Contiguous(sinkOptional, executor);
640 CHECK_RET(sinkOptional != nullptr, ACLNN_ERR_PARAM_NULLPTR);657 CHECK_RET(sinkOptional != nullptr, ACLNN_ERR_PARAM_NULLPTR);
641 }658 }
659+ if (dScaleQOptional) {
660+ dScaleQOptional = l0op::Contiguous(dScaleQOptional, executor);
661+ CHECK_RET(dScaleQOptional != nullptr, ACLNN_ERR_INNER_NULLPTR);
662+ }
663+ if (dScaleKOptional) {
664+ dScaleKOptional = l0op::Contiguous(dScaleKOptional, executor);
665+ CHECK_RET(dScaleKOptional != nullptr, ACLNN_ERR_INNER_NULLPTR);
666+ }
667+ if (dScaleVOptional) {
668+ dScaleVOptional = l0op::Contiguous(dScaleVOptional, executor);
669+ CHECK_RET(dScaleVOptional != nullptr, ACLNN_ERR_INNER_NULLPTR);
670+ }
642 return ACLNN_SUCCESS;671 return ACLNN_SUCCESS;
643}672}
644 673 
645static aclnnStatus PreprocessQKV(const aclTensor *&query, const aclTensor *&key, const aclTensor *&value,674static aclnnStatus PreprocessQKV(const aclTensor *&query, const aclTensor *&key, const aclTensor *&value,
646 const struct FaShapeInfo &shapeInfo, aclOpExecutor *executor)675 const struct FaShapeInfo &shapeInfo, aclOpExecutor *executor)
647{676{
677+ if (!StrideLimited()) {
678+ return ACLNN_SUCCESS;
679+ }
648 if (shapeInfo.needReshape) {680 if (shapeInfo.needReshape) {
649 query = l0op::Reshape(681 query = l0op::Reshape(
650 query, executor->AllocIntArray(shapeInfo.reshapedQueryShape.data(), shapeInfo.reshapedQueryShape.size()),682 query, executor->AllocIntArray(shapeInfo.reshapedQueryShape.data(), shapeInfo.reshapedQueryShape.size()),
@@ -714,6 +746,9 @@ static aclnnStatus PreprocessQKV(const aclTensor *&query, const aclTensor *&key,
714static aclnnStatus Postprocess(const aclTensor *&l0AttentionOutOut, const aclTensor *attentionOutOut,746static aclnnStatus Postprocess(const aclTensor *&l0AttentionOutOut, const aclTensor *attentionOutOut,
715 struct FaShapeInfo &shapeInfo, aclOpExecutor *executor)747 struct FaShapeInfo &shapeInfo, aclOpExecutor *executor)
716{748{
749+ if (!StrideLimited()) {
750+ return ACLNN_SUCCESS;
751+ }
717 if (shapeInfo.inputLayout == InputLayout::SBH && shapeInfo.needPad && !shapeInfo.needTranspose) {752 if (shapeInfo.inputLayout == InputLayout::SBH && shapeInfo.needPad && !shapeInfo.needTranspose) {
718 // (S,B,Hp) -> (S,B,N,Dp)753 // (S,B,Hp) -> (S,B,N,Dp)
719 FVector<int64_t, DIM_NUM_4> paddedSBNDShape{shapeInfo.axes.s1, shapeInfo.axes.b, shapeInfo.axes.n1,754 FVector<int64_t, DIM_NUM_4> paddedSBNDShape{shapeInfo.axes.s1, shapeInfo.axes.b, shapeInfo.axes.n1,
@@ -873,16 +908,20 @@ aclnnStatus aclnnFlashAttentionScoreGetWorkspaceSize(
873 const aclTensor *keyRope = nullptr;908 const aclTensor *keyRope = nullptr;
874 909
875 //检查format是否符合要求910 //检查format是否符合要求
876- if (GetCurrentPlatformInfo().GetSocVersion() != SocVersion::ASCEND910_95) {911+ if (StrideLimited()) {
877 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, 912 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
878 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);913 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);
879 }914 }
880- CHECK_RET(InputDtypeCheck(query, key, value, realShiftOptional, PSE_TYPE_V1, sinkOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);915+ CHECK_RET(InputDtypeCheck(query, key, value, attentionOutOut, realShiftOptional, PSE_TYPE_V1, sinkOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
881 FaShapeInfo shapeInfo;916 FaShapeInfo shapeInfo;
882 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);917 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
883 918 
884 aclOpExecutor *l0Executor = uniqueExecutor.get();919 aclOpExecutor *l0Executor = uniqueExecutor.get();
885- CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, queryRope, keyRope, sinkOptional,920+ const aclTensor *dScaleQ = nullptr;
921+ const aclTensor *dScaleK = nullptr;
922+ const aclTensor *dScaleV = nullptr;
923+ CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional,
924+ attenMaskOptional, queryRope, keyRope, sinkOptional, dScaleQ, dScaleK, dScaleV,
886 l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);925 l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
887 926 
888 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);927 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
@@ -899,7 +938,12 @@ aclnnStatus aclnnFlashAttentionScoreGetWorkspaceSize(
899 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];938 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];
900 // l0SoftmaxOutOut not used now939 // l0SoftmaxOutOut not used now
901 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[DIM_NUM_3];940 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[DIM_NUM_3];
902- 941+ if (l0SoftmaxMaxOut == nullptr || l0SoftmaxSumOut == nullptr || l0AttentionOutOut == nullptr) {
942+ OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "l0SoftmaxMaxOut or l0SoftmaxSumOut or l0AttentionOutOut is null");
943+ *workspaceSize = 0;
944+ uniqueExecutor.ReleaseTo(executor);
945+ return ACLNN_ERR_PARAM_NULLPTR;
946+ }
903 CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,947 CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,
904 ACLNN_ERR_PARAM_NULLPTR);948 ACLNN_ERR_PARAM_NULLPTR);
905 949 
@@ -962,18 +1006,22 @@ aclnnStatus aclnnFlashAttentionVarLenScoreGetWorkspaceSize(
962 const aclTensor *keyRope = nullptr;1006 const aclTensor *keyRope = nullptr;
963 1007
964 //检查format是否符合要求1008 //检查format是否符合要求
965- if (GetCurrentPlatformInfo().GetSocVersion() != SocVersion::ASCEND910_95) {1009+ if (StrideLimited()) {
966 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, 1010 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
967 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);1011 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);
968 }1012 }
969 1013 
970- CHECK_RET(InputDtypeCheck(query, key, value, realShiftOptional, PSE_TYPE_V1, sinkOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);1014+ CHECK_RET(InputDtypeCheck(query, key, value, attentionOutOut, realShiftOptional, PSE_TYPE_V1, sinkOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
971 FaShapeInfo shapeInfo;1015 FaShapeInfo shapeInfo;
972 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo, actualSeqQLenOptional,1016 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo, actualSeqQLenOptional,
973 actualSeqKvLenOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);1017 actualSeqKvLenOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
974 1018 
975 aclOpExecutor *l0Executor = uniqueExecutor.get();1019 aclOpExecutor *l0Executor = uniqueExecutor.get();
976- CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, queryRope, keyRope, sinkOptional,1020+ const aclTensor *dScaleQ = nullptr;
1021+ const aclTensor *dScaleK = nullptr;
1022+ const aclTensor *dScaleV = nullptr;
1023+ CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
1024+ queryRope, keyRope, sinkOptional, dScaleQ, dScaleK, dScaleV,
977 l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);1025 l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
978 1026 
979 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);1027 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
@@ -988,15 +1036,15 @@ aclnnStatus aclnnFlashAttentionVarLenScoreGetWorkspaceSize(
988 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];1036 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];
989 // l0SoftmaxOutOut not used now1037 // l0SoftmaxOutOut not used now
990 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];1038 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];
991- 
992- CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,
993- ACLNN_ERR_PARAM_NULLPTR);
994 if (l0SoftmaxMaxOut == nullptr || l0SoftmaxSumOut == nullptr || l0AttentionOutOut == nullptr) {1039 if (l0SoftmaxMaxOut == nullptr || l0SoftmaxSumOut == nullptr || l0AttentionOutOut == nullptr) {
995 OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "l0SoftmaxMaxOut or l0SoftmaxSumOut or l0AttentionOutOut is null");1040 OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "l0SoftmaxMaxOut or l0SoftmaxSumOut or l0AttentionOutOut is null");
996 *workspaceSize = 0;1041 *workspaceSize = 0;
997 uniqueExecutor.ReleaseTo(executor);1042 uniqueExecutor.ReleaseTo(executor);
998 return ACLNN_ERR_PARAM_NULLPTR;1043 return ACLNN_ERR_PARAM_NULLPTR;
999 }1044 }
1045+ CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,
1046+ ACLNN_ERR_PARAM_NULLPTR);
1047+ 
1000 auto viewCopyResult0 = l0op::ViewCopy(l0SoftmaxMaxOut, softmaxMaxOut, l0Executor);1048 auto viewCopyResult0 = l0op::ViewCopy(l0SoftmaxMaxOut, softmaxMaxOut, l0Executor);
1001 CHECK_RET(viewCopyResult0 != nullptr, ACLNN_ERR_PARAM_NULLPTR);1049 CHECK_RET(viewCopyResult0 != nullptr, ACLNN_ERR_PARAM_NULLPTR);
1002 auto viewCopyResult1 = l0op::ViewCopy(l0SoftmaxSumOut, softmaxSumOut, l0Executor);1050 auto viewCopyResult1 = l0op::ViewCopy(l0SoftmaxSumOut, softmaxSumOut, l0Executor);
@@ -1050,18 +1098,22 @@ aclnnStatus aclnnFlashAttentionScoreV2GetWorkspaceSize(
1050 const aclTensor *keyRope = nullptr;1098 const aclTensor *keyRope = nullptr;
1051 1099
1052 //检查format是否符合要求1100 //检查format是否符合要求
1053- if (GetCurrentPlatformInfo().GetSocVersion() != SocVersion::ASCEND910_95) {1101+ if (StrideLimited()) {
1054 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, 1102 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
1055 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);1103 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);
1056 }1104 }
1057 1105 
1058- CHECK_RET(InputDtypeCheck(query, key, value, realShiftOptional, pseType, nullptr) == ACLNN_SUCCESS,1106+ CHECK_RET(InputDtypeCheck(query, key, value, attentionOutOut, realShiftOptional, pseType, nullptr) == ACLNN_SUCCESS,
1059 ACLNN_ERR_PARAM_INVALID);1107 ACLNN_ERR_PARAM_INVALID);
1060 FaShapeInfo shapeInfo;1108 FaShapeInfo shapeInfo;
1061 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);1109 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
1062 1110 
1063 aclOpExecutor *l0Executor = uniqueExecutor.get();1111 aclOpExecutor *l0Executor = uniqueExecutor.get();
1064- CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, queryRope, keyRope, sinkOptional,1112+ const aclTensor *dScaleQ = nullptr;
1113+ const aclTensor *dScaleK = nullptr;
1114+ const aclTensor *dScaleV = nullptr;
1115+ CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
1116+ queryRope, keyRope, sinkOptional, dScaleQ, dScaleK, dScaleV,
1065 l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);1117 l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
1066 1118 
1067 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);1119 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
@@ -1076,7 +1128,12 @@ aclnnStatus aclnnFlashAttentionScoreV2GetWorkspaceSize(
1076 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];1128 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];
1077 // l0SoftmaxOutOut not used now1129 // l0SoftmaxOutOut not used now
1078 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];1130 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];
1079- 1131+ if (l0SoftmaxMaxOut == nullptr || l0SoftmaxSumOut == nullptr || l0AttentionOutOut == nullptr) {
1132+ OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "l0SoftmaxMaxOut or l0SoftmaxSumOut or l0AttentionOutOut is null");
1133+ *workspaceSize = 0;
1134+ uniqueExecutor.ReleaseTo(executor);
1135+ return ACLNN_ERR_PARAM_NULLPTR;
1136+ }
1080 CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,1137 CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,
1081 ACLNN_ERR_PARAM_NULLPTR);1138 ACLNN_ERR_PARAM_NULLPTR);
1082 1139 
@@ -1133,18 +1190,22 @@ aclnnStatus aclnnFlashAttentionScoreV3GetWorkspaceSize(
1133 const aclTensor *keyRope = nullptr;1190 const aclTensor *keyRope = nullptr;
1134 1191
1135 //检查format是否符合要求1192 //检查format是否符合要求
1136- if (GetCurrentPlatformInfo().GetSocVersion() != SocVersion::ASCEND910_95) {1193+ if (StrideLimited()) {
1137 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, 1194 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
1138 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);1195 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);
1139 }1196 }
1140 1197 
1141- CHECK_RET(InputDtypeCheck(query, key, value, realShiftOptional, pseType, sinkOptional) == ACLNN_SUCCESS,1198+ CHECK_RET(InputDtypeCheck(query, key, value, attentionOutOut, realShiftOptional, pseType, sinkOptional) == ACLNN_SUCCESS,
1142 ACLNN_ERR_PARAM_INVALID);1199 ACLNN_ERR_PARAM_INVALID);
1143 FaShapeInfo shapeInfo;1200 FaShapeInfo shapeInfo;
1144 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);1201 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
1145 1202 
1146 aclOpExecutor *l0Executor = uniqueExecutor.get();1203 aclOpExecutor *l0Executor = uniqueExecutor.get();
1147- CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, queryRope, keyRope, sinkOptional,1204+ const aclTensor *dScaleQ = nullptr;
1205+ const aclTensor *dScaleK = nullptr;
1206+ const aclTensor *dScaleV = nullptr;
1207+ CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
1208+ queryRope, keyRope, sinkOptional, dScaleQ, dScaleK, dScaleV,
1148 l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);1209 l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
1149 1210 
1150 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);1211 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
@@ -1164,7 +1225,12 @@ aclnnStatus aclnnFlashAttentionScoreV3GetWorkspaceSize(
1164 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];1225 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];
1165 // l0SoftmaxOutOut not used now1226 // l0SoftmaxOutOut not used now
1166 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];1227 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];
1167- 1228+ if (l0SoftmaxMaxOut == nullptr || l0SoftmaxSumOut == nullptr || l0AttentionOutOut == nullptr) {
1229+ OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "l0SoftmaxMaxOut or l0SoftmaxSumOut or l0AttentionOutOut is null");
1230+ *workspaceSize = 0;
1231+ uniqueExecutor.ReleaseTo(executor);
1232+ return ACLNN_ERR_PARAM_NULLPTR;
1233+ }
1168 CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,1234 CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,
1169 ACLNN_ERR_PARAM_NULLPTR);1235 ACLNN_ERR_PARAM_NULLPTR);
1170 1236 
@@ -1189,6 +1255,100 @@ aclnnStatus aclnnFlashAttentionScoreV3(void *workspace, uint64_t workspaceSize,
1189 return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);1255 return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);
1190}1256}
1191 1257 
1258+aclnnStatus aclnnFlashAttentionScoreV4GetWorkspaceSize(
1259+ const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *realShiftOptional,
1260+ const aclTensor *dropMaskOptional, const aclTensor *paddingMaskOptional, const aclTensor *attenMaskOptional,
1261+ const aclTensor *queryRopeOptional, const aclTensor *keyRopeOptional, const aclTensor *dScaleQOptional,
1262+ const aclTensor *dScaleKOptional, const aclTensor *dScaleVOptional,
1263+ const aclTensor *sinkOptional, const aclIntArray *prefixOptional,
1264+ const aclIntArray *actualSeqQLenOptional, const aclIntArray *actualSeqKvLenOptional,
1265+ const aclIntArray *qStartIdxOptional, const aclIntArray *kvStartIdxOptional, double scaleValue,
1266+ double keepProb, int64_t preTokens, int64_t nextTokens, int64_t headNum, char *inputLayout, int64_t innerPrecise,
1267+ int64_t sparseMode, int64_t outDtype, int64_t pseType, char *softmaxOutLayout, int64_t seed, int64_t offset,
1268+ const aclTensor *softmaxMaxOut, const aclTensor *softmaxSumOut, const aclTensor *softmaxOutOut,
1269+ const aclTensor *attentionOutOut, uint64_t *workspaceSize, aclOpExecutor **executor)
1270+{
1271+ CHECK_RET(CheckFaParam(query, key, value, inputLayout, softmaxMaxOut, softmaxSumOut, attentionOutOut,
1272+ workspaceSize, executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
1273+ L2_DFX_PHASE_1(aclnnFlashAttentionScoreV4,
1274+ DFX_IN(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional,
1275+ attenMaskOptional, queryRopeOptional, keyRopeOptional, dScaleQOptional, dScaleKOptional,
1276+ dScaleVOptional, sinkOptional, prefixOptional, actualSeqQLenOptional, actualSeqKvLenOptional,
1277+ qStartIdxOptional, kvStartIdxOptional, scaleValue, keepProb, preTokens, nextTokens, headNum,
1278+ inputLayout, innerPrecise, sparseMode, outDtype, pseType, seed, offset),
1279+ DFX_OUT(softmaxMaxOut, softmaxSumOut, softmaxOutOut, attentionOutOut));
1280+ 
1281+ auto uniqueExecutor = CREATE_EXECUTOR();
1282+ CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);
1283+ // b, n1, s1 为0时，不进行任何处理
1284+ // n2, s2, d 为0时，直接调用l0接口处理
1285+ if (softmaxMaxOut->IsEmpty() && softmaxSumOut->IsEmpty() && attentionOutOut->IsEmpty()) {
1286+ *workspaceSize = 0;
1287+ uniqueExecutor.ReleaseTo(executor);
1288+ return ACLNN_SUCCESS;
1289+ }
1290+ CHECK_RET(InputDtypeCheck(query, key, value, attentionOutOut, realShiftOptional, pseType, sinkOptional) ==
1291+ ACLNN_SUCCESS,
1292+ ACLNN_ERR_PARAM_INVALID);
1293+ FaShapeInfo shapeInfo;
1294+ CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo) == ACLNN_SUCCESS,
1295+ ACLNN_ERR_PARAM_INVALID);
1296+ 
1297+ aclOpExecutor *l0Executor = uniqueExecutor.get();
1298+ 
1299+ CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
1300+ queryRopeOptional, keyRopeOptional, sinkOptional, dScaleQOptional, dScaleKOptional,
1301+ dScaleVOptional, l0Executor) == ACLNN_SUCCESS,
1302+ ACLNN_ERR_INNER_NULLPTR);
1303+ 
1304+ CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
1305+ 
1306+ // sink shape is 0
1307+ if (sinkOptional != nullptr && sinkOptional->GetViewShape().GetDimNum() == 1 &&
1308+ sinkOptional->GetViewShape()[0] == 0) {
1309+ sinkOptional = nullptr;
1310+ }
1311+ auto l0FlashAttentionScoreOuts = l0op::FlashAttentionScore(
1312+ query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, sinkOptional,
1313+ prefixOptional, actualSeqQLenOptional, actualSeqKvLenOptional, qStartIdxOptional, kvStartIdxOptional,
1314+ dScaleQOptional, dScaleKOptional, dScaleVOptional, queryRopeOptional, keyRopeOptional,
1315+ scaleValue, keepProb, preTokens, nextTokens, headNum, shapeInfo.l0InputLayoutStr.c_str(),
1316+ innerPrecise, sparseMode, pseType, seed, offset, outDtype, softmaxOutLayout, l0Executor);
1317+ 
1318+ auto l0SoftmaxMaxOut = l0FlashAttentionScoreOuts[0];
1319+ auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];
1320+ // l0SoftmaxOutOut not used now
1321+ auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];
1322+ if (l0SoftmaxMaxOut == nullptr || l0SoftmaxSumOut == nullptr || l0AttentionOutOut == nullptr) {
1323+ OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "l0SoftmaxMaxOut or l0SoftmaxSumOut or l0AttentionOutOut is null");
1324+ *workspaceSize = 0;
1325+ uniqueExecutor.ReleaseTo(executor);
1326+ return ACLNN_ERR_PARAM_NULLPTR;
1327+ }
1328+ CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,
1329+ ACLNN_ERR_PARAM_NULLPTR);
1330+ 
1331+ auto viewCopyResult0 = l0op::ViewCopy(l0SoftmaxMaxOut, softmaxMaxOut, l0Executor);
1332+ CHECK_RET(viewCopyResult0 != nullptr, ACLNN_ERR_PARAM_NULLPTR);
1333+ auto viewCopyResult1 = l0op::ViewCopy(l0SoftmaxSumOut, softmaxSumOut, l0Executor);
1334+ CHECK_RET(viewCopyResult1 != nullptr, ACLNN_ERR_PARAM_NULLPTR);
1335+ // l0SoftmaxOutOut not used now
1336+ auto viewCopyResult3 = l0op::ViewCopy(l0AttentionOutOut, attentionOutOut, l0Executor);
1337+ CHECK_RET(viewCopyResult3 != nullptr, ACLNN_ERR_PARAM_NULLPTR);
1338+ 
1339+ *workspaceSize = uniqueExecutor->GetWorkspaceSize();
1340+ uniqueExecutor.ReleaseTo(executor);
1341+ return ACLNN_SUCCESS;
1342+}
1343+ 
1344+aclnnStatus aclnnFlashAttentionScoreV4(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor,
1345+ const aclrtStream stream)
1346+{
1347+ L2_DFX_PHASE_2(aclnnFlashAttentionScoreV4);
1348+ // 固定写法，调用框架能力，完成计算
1349+ return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);
1350+}
1351+ 
1192aclnnStatus aclnnFlashAttentionVarLenScoreV2GetWorkspaceSize(1352aclnnStatus aclnnFlashAttentionVarLenScoreV2GetWorkspaceSize(
1193 const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *realShiftOptional,1353 const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *realShiftOptional,
1194 const aclTensor *dropMaskOptional, const aclTensor *paddingMaskOptional, const aclTensor *attenMaskOptional,1354 const aclTensor *dropMaskOptional, const aclTensor *paddingMaskOptional, const aclTensor *attenMaskOptional,
@@ -1230,18 +1390,22 @@ aclnnStatus aclnnFlashAttentionVarLenScoreV2GetWorkspaceSize(
1230 const aclTensor *keyRope = nullptr;1390 const aclTensor *keyRope = nullptr;
1231 1391
1232 //检查format是否符合要求1392 //检查format是否符合要求
1233- if (GetCurrentPlatformInfo().GetSocVersion() != SocVersion::ASCEND910_95) {1393+ if (StrideLimited()) {
1234 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, 1394 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
1235 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);1395 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);
1236 }1396 }
1237 1397 
1238- CHECK_RET(InputDtypeCheck(query, key, value, realShiftOptional, pseType, sinkOptional) == ACLNN_SUCCESS,1398+ CHECK_RET(InputDtypeCheck(query, key, value, attentionOutOut, realShiftOptional, pseType, sinkOptional) == ACLNN_SUCCESS,
1239 ACLNN_ERR_PARAM_INVALID);1399 ACLNN_ERR_PARAM_INVALID);
1240 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo, actualSeqQLenOptional,1400 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo, actualSeqQLenOptional,
1241 actualSeqKvLenOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);1401 actualSeqKvLenOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
1242 1402 
1243 aclOpExecutor *l0Executor = uniqueExecutor.get();1403 aclOpExecutor *l0Executor = uniqueExecutor.get();
1244- CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, queryRope, keyRope, sinkOptional,1404+ const aclTensor *dScaleQ = nullptr;
1405+ const aclTensor *dScaleK = nullptr;
1406+ const aclTensor *dScaleV = nullptr;
1407+ CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
1408+ queryRope, keyRope, sinkOptional, dScaleQ, dScaleK, dScaleV,
1245 l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);1409 l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
1246 1410 
1247 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);1411 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
@@ -1256,15 +1420,15 @@ aclnnStatus aclnnFlashAttentionVarLenScoreV2GetWorkspaceSize(
1256 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];1420 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];
1257 // l0SoftmaxOutOut not used now1421 // l0SoftmaxOutOut not used now
1258 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];1422 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];
1259- 1423+ if (l0SoftmaxMaxOut == nullptr || l0SoftmaxSumOut == nullptr || l0AttentionOutOut == nullptr) {
1424+ OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "l0SoftmaxMaxOut or l0SoftmaxSumOut or l0AttentionOutOut is null");
1425+ *workspaceSize = 0;
1426+ uniqueExecutor.ReleaseTo(executor);
1427+ return ACLNN_ERR_PARAM_NULLPTR;
1428+ }
1260 CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,1429 CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,
1261 ACLNN_ERR_PARAM_NULLPTR);1430 ACLNN_ERR_PARAM_NULLPTR);
1262- if (l0SoftmaxMaxOut == nullptr || l0SoftmaxSumOut == nullptr || l0AttentionOutOut == nullptr) {1431+ 
1263- OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "l0SoftmaxMaxOut or l0SoftmaxSumOut or l0AttentionOutOut is null");
1264- *workspaceSize = 0;
1265- uniqueExecutor.ReleaseTo(executor);
1266- return ACLNN_ERR_PARAM_NULLPTR;
1267- }
1268 auto viewCopyResult0 = l0op::ViewCopy(l0SoftmaxMaxOut, softmaxMaxOut, l0Executor);1432 auto viewCopyResult0 = l0op::ViewCopy(l0SoftmaxMaxOut, softmaxMaxOut, l0Executor);
1269 CHECK_RET(viewCopyResult0 != nullptr, ACLNN_ERR_PARAM_NULLPTR);1433 CHECK_RET(viewCopyResult0 != nullptr, ACLNN_ERR_PARAM_NULLPTR);
1270 auto viewCopyResult1 = l0op::ViewCopy(l0SoftmaxSumOut, softmaxSumOut, l0Executor);1434 auto viewCopyResult1 = l0op::ViewCopy(l0SoftmaxSumOut, softmaxSumOut, l0Executor);
@@ -1326,20 +1490,24 @@ aclnnStatus aclnnFlashAttentionVarLenScoreV3GetWorkspaceSize(
1326 const aclTensor *sinkOptional = nullptr;1490 const aclTensor *sinkOptional = nullptr;
1327 1491
1328 //检查format是否符合要求1492 //检查format是否符合要求
1329- if (GetCurrentPlatformInfo().GetSocVersion() != SocVersion::ASCEND910_95) {1493+ if (StrideLimited()) {
1330 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, 1494 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
1331 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);1495 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);
1332 }1496 }
1333 1497 
1334- CHECK_RET(InputDtypeCheck(query, key, value, realShiftOptional, pseType, sinkOptional) == ACLNN_SUCCESS,1498+ CHECK_RET(InputDtypeCheck(query, key, value, attentionOutOut, realShiftOptional, pseType, sinkOptional) == ACLNN_SUCCESS,
1335 ACLNN_ERR_PARAM_INVALID);1499 ACLNN_ERR_PARAM_INVALID);
1336 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo, actualSeqQLenOptional,1500 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo, actualSeqQLenOptional,
1337 actualSeqKvLenOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);1501 actualSeqKvLenOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
1338 1502 
1339 aclOpExecutor *l0Executor = uniqueExecutor.get();1503 aclOpExecutor *l0Executor = uniqueExecutor.get();
1340 1504 
1341- CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, queryRope, keyRope, sinkOptional,1505+ const aclTensor *dScaleQ = nullptr;
1342- l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);1506+ const aclTensor *dScaleK = nullptr;
1507+ const aclTensor *dScaleV = nullptr;
1508+ CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
1509+ queryRope, keyRope, sinkOptional, dScaleQ, dScaleK, dScaleV,
1510+ l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
1343 1511 
1344 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);1512 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
1345 1513 
@@ -1356,15 +1524,15 @@ aclnnStatus aclnnFlashAttentionVarLenScoreV3GetWorkspaceSize(
1356 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];1524 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];
1357 // l0SoftmaxOutOut not used now1525 // l0SoftmaxOutOut not used now
1358 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];1526 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];
1359- 1527+ if (l0SoftmaxMaxOut == nullptr || l0SoftmaxSumOut == nullptr || l0AttentionOutOut == nullptr) {
1528+ OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "l0SoftmaxMaxOut or l0SoftmaxSumOut or l0AttentionOutOut is null");
1529+ *workspaceSize = 0;
1530+ uniqueExecutor.ReleaseTo(executor);
1531+ return ACLNN_ERR_PARAM_NULLPTR;
1532+ }
1360 CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,1533 CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,
1361 ACLNN_ERR_PARAM_NULLPTR);1534 ACLNN_ERR_PARAM_NULLPTR);
1362- if (l0SoftmaxMaxOut == nullptr || l0SoftmaxSumOut == nullptr || l0AttentionOutOut == nullptr) {1535+ 
1363- OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "l0SoftmaxMaxOut or l0SoftmaxSumOut or l0AttentionOutOut is null");
1364- *workspaceSize = 0;
1365- uniqueExecutor.ReleaseTo(executor);
1366- return ACLNN_ERR_PARAM_NULLPTR;
1367- }
1368 auto viewCopyResult0 = l0op::ViewCopy(l0SoftmaxMaxOut, softmaxMaxOut, l0Executor);1536 auto viewCopyResult0 = l0op::ViewCopy(l0SoftmaxMaxOut, softmaxMaxOut, l0Executor);
1369 CHECK_RET(viewCopyResult0 != nullptr, ACLNN_ERR_PARAM_NULLPTR);1537 CHECK_RET(viewCopyResult0 != nullptr, ACLNN_ERR_PARAM_NULLPTR);
1370 auto viewCopyResult1 = l0op::ViewCopy(l0SoftmaxSumOut, softmaxSumOut, l0Executor);1538 auto viewCopyResult1 = l0op::ViewCopy(l0SoftmaxSumOut, softmaxSumOut, l0Executor);
@@ -1424,18 +1592,22 @@ aclnnStatus aclnnFlashAttentionVarLenScoreV4GetWorkspaceSize(
1424 const aclTensor *keyRope = nullptr;1592 const aclTensor *keyRope = nullptr;
1425 1593
1426 //检查format是否符合要求1594 //检查format是否符合要求
1427- if (GetCurrentPlatformInfo().GetSocVersion() != SocVersion::ASCEND910_95) {1595+ if (StrideLimited()) {
1428 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, 1596 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
1429 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);1597 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);
1430 }1598 }
1431 1599 
1432- CHECK_RET(InputDtypeCheck(query, key, value, realShiftOptional, PSE_TYPE_V1, sinkOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);1600+ CHECK_RET(InputDtypeCheck(query, key, value, attentionOutOut, realShiftOptional, PSE_TYPE_V1, sinkOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
1433 FaShapeInfo shapeInfo;1601 FaShapeInfo shapeInfo;
1434 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo, actualSeqQLenOptional,1602 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo, actualSeqQLenOptional,
1435 actualSeqKvLenOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);1603 actualSeqKvLenOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
1436 1604 
1437 aclOpExecutor *l0Executor = uniqueExecutor.get();1605 aclOpExecutor *l0Executor = uniqueExecutor.get();
1438- CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, queryRope, keyRope, sinkOptional,1606+ const aclTensor *dScaleQ = nullptr;
1607+ const aclTensor *dScaleK = nullptr;
1608+ const aclTensor *dScaleV = nullptr;
1609+ CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
1610+ queryRope, keyRope, sinkOptional, dScaleQ, dScaleK, dScaleV,
1439 l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);1611 l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
1440 1612 
1441 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);1613 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
@@ -1450,15 +1622,15 @@ aclnnStatus aclnnFlashAttentionVarLenScoreV4GetWorkspaceSize(
1450 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];1622 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];
1451 // l0SoftmaxOutOut not used now1623 // l0SoftmaxOutOut not used now
1452 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];1624 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];
1453- 
1454- CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,
1455- ACLNN_ERR_PARAM_NULLPTR);
1456 if (l0SoftmaxMaxOut == nullptr || l0SoftmaxSumOut == nullptr || l0AttentionOutOut == nullptr) {1625 if (l0SoftmaxMaxOut == nullptr || l0SoftmaxSumOut == nullptr || l0AttentionOutOut == nullptr) {
1457 OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "l0SoftmaxMaxOut or l0SoftmaxSumOut or l0AttentionOutOut is null");1626 OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "l0SoftmaxMaxOut or l0SoftmaxSumOut or l0AttentionOutOut is null");
1458 *workspaceSize = 0;1627 *workspaceSize = 0;
1459 uniqueExecutor.ReleaseTo(executor);1628 uniqueExecutor.ReleaseTo(executor);
1460 return ACLNN_ERR_PARAM_NULLPTR;1629 return ACLNN_ERR_PARAM_NULLPTR;
1461 }1630 }
1631+ CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,
1632+ ACLNN_ERR_PARAM_NULLPTR);
1633+ 
1462 auto viewCopyResult0 = l0op::ViewCopy(l0SoftmaxMaxOut, softmaxMaxOut, l0Executor);1634 auto viewCopyResult0 = l0op::ViewCopy(l0SoftmaxMaxOut, softmaxMaxOut, l0Executor);
1463 CHECK_RET(viewCopyResult0 != nullptr, ACLNN_ERR_PARAM_NULLPTR);1635 CHECK_RET(viewCopyResult0 != nullptr, ACLNN_ERR_PARAM_NULLPTR);
1464 auto viewCopyResult1 = l0op::ViewCopy(l0SoftmaxSumOut, softmaxSumOut, l0Executor);1636 auto viewCopyResult1 = l0op::ViewCopy(l0SoftmaxSumOut, softmaxSumOut, l0Executor);
@@ -1517,21 +1689,25 @@ aclnnStatus aclnnFlashAttentionVarLenScoreV5GetWorkspaceSize(
1517 }1689 }
1518 1690 
1519 //检查format是否符合要求1691 //检查format是否符合要求
1520- if (GetCurrentPlatformInfo().GetSocVersion() != SocVersion::ASCEND910_95) {1692+ if (StrideLimited()) {
1521 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, 1693 CHECK_RET(CheckFormat(query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
1522 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);1694 sinkOptional, softmaxMaxOut, softmaxSumOut, attentionOutOut), ACLNN_ERR_PARAM_INVALID);
1523 }1695 }
1524 1696 
1525 FaShapeInfo shapeInfo;1697 FaShapeInfo shapeInfo;
1526- CHECK_RET(InputDtypeCheck(query, key, value, realShiftOptional, pseType, sinkOptional) == ACLNN_SUCCESS,1698+ CHECK_RET(InputDtypeCheck(query, key, value, attentionOutOut, realShiftOptional, pseType, sinkOptional) == ACLNN_SUCCESS,
1527 ACLNN_ERR_PARAM_INVALID);1699 ACLNN_ERR_PARAM_INVALID);
1528 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo, actualSeqQLenOptional,1700 CHECK_RET(AnalysisInput(query, key, value, inputLayout, headNum, shapeInfo, actualSeqQLenOptional,
1529 actualSeqKvLenOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);1701 actualSeqKvLenOptional) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
1530 1702 
1531 aclOpExecutor *l0Executor = uniqueExecutor.get();1703 aclOpExecutor *l0Executor = uniqueExecutor.get();
1532 1704 
1533- CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, queryRope, keyRope, sinkOptional,1705+ const aclTensor *dScaleQ = nullptr;
1534- l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);1706+ const aclTensor *dScaleK = nullptr;
1707+ const aclTensor *dScaleV = nullptr;
1708+ CHECK_RET(Contiguous(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
1709+ queryRope, keyRope, sinkOptional, dScaleQ, dScaleK, dScaleV,
1710+ l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
1535 1711 
1536 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);1712 CHECK_RET(PreprocessQKV(query, key, value, shapeInfo, l0Executor) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
1537 1713 
@@ -1552,15 +1728,15 @@ aclnnStatus aclnnFlashAttentionVarLenScoreV5GetWorkspaceSize(
1552 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];1728 auto l0SoftmaxSumOut = l0FlashAttentionScoreOuts[1];
1553 // l0SoftmaxOutOut not used now1729 // l0SoftmaxOutOut not used now
1554 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];1730 auto l0AttentionOutOut = l0FlashAttentionScoreOuts[3];
1555- 1731+ if (l0SoftmaxMaxOut == nullptr || l0SoftmaxSumOut == nullptr || l0AttentionOutOut == nullptr) {
1732+ OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "l0SoftmaxMaxOut or l0SoftmaxSumOut or l0AttentionOutOut is null");
1733+ *workspaceSize = 0;
1734+ uniqueExecutor.ReleaseTo(executor);
1735+ return ACLNN_ERR_PARAM_NULLPTR;
1736+ }
1556 CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,1737 CHECK_RET(Postprocess(l0AttentionOutOut, attentionOutOut, shapeInfo, l0Executor) == ACLNN_SUCCESS,
1557 ACLNN_ERR_PARAM_NULLPTR);1738 ACLNN_ERR_PARAM_NULLPTR);
1558- if (l0SoftmaxMaxOut == nullptr || l0SoftmaxSumOut == nullptr || l0AttentionOutOut == nullptr) {1739+ 
1559- OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "l0SoftmaxMaxOut or l0SoftmaxSumOut or l0AttentionOutOut is null");
1560- *workspaceSize = 0;
1561- uniqueExecutor.ReleaseTo(executor);
1562- return ACLNN_ERR_PARAM_NULLPTR;
1563- }
1564 auto viewCopyResult0 = l0op::ViewCopy(l0SoftmaxMaxOut, softmaxMaxOut, l0Executor);1740 auto viewCopyResult0 = l0op::ViewCopy(l0SoftmaxMaxOut, softmaxMaxOut, l0Executor);
1565 CHECK_RET(viewCopyResult0 != nullptr, ACLNN_ERR_PARAM_NULLPTR);1741 CHECK_RET(viewCopyResult0 != nullptr, ACLNN_ERR_PARAM_NULLPTR);
1566 auto viewCopyResult1 = l0op::ViewCopy(l0SoftmaxSumOut, softmaxSumOut, l0Executor);1742 auto viewCopyResult1 = l0op::ViewCopy(l0SoftmaxSumOut, softmaxSumOut, l0Executor);
Mattention/flash_attention_score/op_api/aclnn_flash_attention_score.h+52-0
@@ -136,6 +136,58 @@ aclnnStatus aclnnFlashAttentionScoreV3(
136 aclOpExecutor *executor,136 aclOpExecutor *executor,
137 const aclrtStream stream);137 const aclrtStream stream);
138 138 
139+/**
140+ * @brief aclnnFlashAttentionScoreV4的第一段接口，根据具体的计算流程，计算workspace大小。
141+ * @domain aclnn_ops_train
142+ */
143+aclnnStatus aclnnFlashAttentionScoreV4GetWorkspaceSize(
CANN-robot
CANN-robotCANN-robot1月20日

API设计一致性: aclnnFlashAttentionScoreV4GetWorkspaceSize函数参数列表与其他版本（如V2、V3）存在不一致性，新增了多个可选参数（queryRopeOptional, keyRopeOptional, dScaleQOptional等），但函数命名和注释没有明确说明这些新增参数的具体用途和版本差异。这可能导致API使用者混淆，不知道何时应该使用哪个版本。

问题类型: API设计一致性 文件路径: attention/flash_attention_score/op_api/aclnn_flash_attention_score.h 行号: 143 问题代码:

aclnnStatus aclnnFlashAttentionScoreV4GetWorkspaceSize(
    const aclTensor *query,
    const aclTensor *key,
    const aclTensor *value,
    const aclTensor *realShiftOptional,
    const aclTensor *dropMaskOptional,
    const aclTensor *paddingMaskOptional,
    const aclTensor *attenMaskOptional,
    const aclTensor *queryRopeOptional,
    const aclTensor *keyRopeOptional,
    const aclTensor *dScaleQOptional,
    const aclTensor *dScaleKOptional,
    const aclTensor *dScaleVOptional,
    const aclTensor *sinkOptional,
    const aclIntArray *prefixOptional,
    const aclIntArray *actualSeqQLenOptional, /*varlen only*/
    const aclIntArray *actualSeqKvLenOptional, /*varlen only*/
    const aclIntArray *qStartIdxOptional,
    const aclIntArray *kvStartIdxOptional,
    double scaleValue,
    double keepProb,
    int64_t preTokens,
    int64_t nextTokens,
    int64_t headNum,
    char *inputLayout,
    int64_t innerPrecise,
    int64_t sparseMode,
    int64_t outDtype,
    int64_t pseType,
    char *softmaxOutLayout,
    int64_t seed, /*dropout ADD*/
    int64_t offset, /*dropout ADD*/
    const aclTensor *softmaxMaxOut,
    const aclTensor *softmaxSumOut,
    const aclTensor *softmaxOutOut,
    const aclTensor *attentionOutOut,
    uint64_t *workspaceSize,
    aclOpExecutor **executor);

修改建议:

1. 在函数注释中明确说明V4版本新增的功能特性（如RoPE支持、动态缩放因子、dropout种子/偏移等）。2. 考虑是否需要在头文件中添加版本功能对比说明，帮助开发者选择合适的API版本。3. 确保所有可选参数在函数实现中有适当的空值检查。

此评论由代码审查工具自动生成

likedislike
144+ const aclTensor *query,
145+ const aclTensor *key,
146+ const aclTensor *value,
147+ const aclTensor *realShiftOptional,
148+ const aclTensor *dropMaskOptional,
149+ const aclTensor *paddingMaskOptional,
150+ const aclTensor *attenMaskOptional,
151+ const aclTensor *queryRopeOptional,
152+ const aclTensor *keyRopeOptional,
153+ const aclTensor *dScaleQOptional,
154+ const aclTensor *dScaleKOptional,
155+ const aclTensor *dScaleVOptional,
156+ const aclTensor *sinkOptional,
157+ const aclIntArray *prefixOptional,
158+ const aclIntArray *actualSeqQLenOptional, /*varlen only*/
CANN-robot
CANN-robotCANN-robot1月20日

代码结构与可维护性: 参数actualSeqQLenOptional和actualSeqKvLenOptional的注释'/varlen only/'表明这些参数仅对变长序列有效，但函数名是FlashAttentionScoreV4（不是VarLen版本）。这种命名和参数的不一致可能导致调用者困惑。

问题类型: 代码结构与可维护性 文件路径: attention/flash_attention_score/op_api/aclnn_flash_attention_score.h 行号: 158 问题代码:

const aclIntArray *actualSeqQLenOptional, /*varlen only*/
const aclIntArray *actualSeqKvLenOptional, /*varlen only*/

修改建议:

1. 重新评估函数命名，如果支持变长序列，考虑更合适的名称；2. 在函数文档中明确说明这些参数的使用条件；3. 如果某些参数只在特定条件下使用，考虑使用条件编译或不同的函数变体。

此评论由代码审查工具自动生成

likedislike
159+ const aclIntArray *actualSeqKvLenOptional, /*varlen only*/
160+ const aclIntArray *qStartIdxOptional,
161+ const aclIntArray *kvStartIdxOptional,
162+ double scaleValue,
163+ double keepProb,
164+ int64_t preTokens,
165+ int64_t nextTokens,
166+ int64_t headNum,
167+ char *inputLayout,
168+ int64_t innerPrecise,
169+ int64_t sparseMode,
170+ int64_t outDtype,
171+ int64_t pseType,
172+ char *softmaxOutLayout,
173+ int64_t seed, /*dropout ADD*/
CANN-robot
CANN-robotCANN-robot1月20日

代码结构与可维护性: 注释中的'/dropout ADD/'表明这是新增的参数，但注释风格不一致。其他参数没有类似的注释说明新增或修改情况，这不利于代码维护和版本追踪。

问题类型: 代码结构与可维护性 文件路径: attention/flash_attention_score/op_api/aclnn_flash_attention_score.h 行号: 173 问题代码:

int64_t seed, /*dropout ADD*/
int64_t offset, /*dropout ADD*/

修改建议:

1. 统一注释风格，对所有新增或修改的参数添加说明；2. 考虑使用版本控制系统来追踪变更，而不是在代码中添加临时注释；3. 如果这是API演进的一部分，应该在头文件顶部添加版本变更说明。

此评论由代码审查工具自动生成

likedislike
174+ int64_t offset, /*dropout ADD*/
175+ const aclTensor *softmaxMaxOut,
176+ const aclTensor *softmaxSumOut,
177+ const aclTensor *softmaxOutOut,
178+ const aclTensor *attentionOutOut,
179+ uint64_t *workspaceSize,
180+ aclOpExecutor **executor);
181+ 
182+/**
183+ * @brief aclnnFlashAttentionScoreV4的第二段接口，用于执行计算。
184+ */
185+aclnnStatus aclnnFlashAttentionScoreV4(
186+ void *workspace,
187+ uint64_t workspaceSize,
188+ aclOpExecutor *executor,
189+ const aclrtStream stream);
190+ 
139 191 
140/**192/**
141 * @brief aclnnFlashAttentionVarLenScoreV2的第一段接口，根据具体的计算流程，计算workspace大小。193 * @brief aclnnFlashAttentionVarLenScoreV2的第一段接口，根据具体的计算流程，计算workspace大小。
Mattention/flash_attention_score/op_api/flash_attention_score.cpp+5-1
@@ -136,7 +136,7 @@ const std::array<const aclTensor *, 4> FlashAttentionScore(
136 return {nullptr, nullptr, nullptr, nullptr};136 return {nullptr, nullptr, nullptr, nullptr};
137 }137 }
138 138 
139- ADD_TO_LAUNCHER_LIST_AICORE(139+ ret = ADD_TO_LAUNCHER_LIST_AICORE(
140 FlashAttentionScore,140 FlashAttentionScore,
141 OP_INPUT(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,141 OP_INPUT(query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional,
142 prefixOptionalTensor, actualSeqQLen, actualSeqKvLen, qStartIdxOptionalTensor, kvStartIdxOptionalTensor,142 prefixOptionalTensor, actualSeqQLen, actualSeqKvLen, qStartIdxOptionalTensor, kvStartIdxOptionalTensor,
@@ -145,6 +145,10 @@ const std::array<const aclTensor *, 4> FlashAttentionScore(
145 OP_ATTR(static_cast<float>(scaleValue), static_cast<float>(keepProb), preTockens,145 OP_ATTR(static_cast<float>(scaleValue), static_cast<float>(keepProb), preTockens,
146 nextTockens, headNum, inputLayout, innerPrecise, sparseMode, pseType,146 nextTockens, headNum, inputLayout, innerPrecise, sparseMode, pseType,
147 seed, offset, outDtype, softmaxOutLayout));147 seed, offset, outDtype, softmaxOutLayout));
148+ if (ret != ACLNN_SUCCESS) {
149+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "FlashAttentionScore launch kernel failed.");
150+ return {nullptr, nullptr, nullptr, nullptr};
151+ }
148 return {softmaxMaxOut, softmaxSumOut, softmaxOutOut, attentionOutOut};152 return {softmaxMaxOut, softmaxSumOut, softmaxOutOut, attentionOutOut};
149}153}
150 154 
Mattention/flash_attention_score/op_host/arch32/flash_attention_score_tiling_general.cpp+6-6
@@ -4533,11 +4533,11 @@ protected:
4533};4533};
4534 4534 
4535// NOTE manually initialize tiling data in hostapi scenario in highest priority template4535// NOTE manually initialize tiling data in hostapi scenario in highest priority template
4536-REGISTER_TILING_TEMPLATE_WITH_SOCVERSION(FlashAttentionScore, FlashAttentionScoreTilingDropMask, std::vector<int32_t>({static_cast<int32_t>(platform_ascendc::SocVersion::ASCEND910B), static_cast<int32_t>(platform_ascendc::SocVersion::ASCEND910_93)}), 90);4536+REGISTER_TILING_TEMPLATE_WITH_ARCH(FlashAttentionScore, FlashAttentionScoreTilingDropMask, std::vector<int32_t>({static_cast<int32_t>(NpuArch::DAV_2201)}), 90);
4537-REGISTER_TILING_TEMPLATE_WITH_SOCVERSION(FlashAttentionScore, FlashAttentionVarLenScoreTiling, std::vector<int32_t>({static_cast<int32_t>(platform_ascendc::SocVersion::ASCEND910B), static_cast<int32_t>(platform_ascendc::SocVersion::ASCEND910_93)}), 94);4537+REGISTER_TILING_TEMPLATE_WITH_ARCH(FlashAttentionScore, FlashAttentionVarLenScoreTiling, std::vector<int32_t>({static_cast<int32_t>(NpuArch::DAV_2201)}), 94);
4538-REGISTER_TILING_TEMPLATE_WITH_SOCVERSION(FlashAttentionScore, FlashAttentionScoreTilingS1s2Bn2gs1SameAB, std::vector<int32_t>({static_cast<int32_t>(platform_ascendc::SocVersion::ASCEND910B), static_cast<int32_t>(platform_ascendc::SocVersion::ASCEND910_93)}), 95);4538+REGISTER_TILING_TEMPLATE_WITH_ARCH(FlashAttentionScore, FlashAttentionScoreTilingS1s2Bn2gs1SameAB, std::vector<int32_t>({static_cast<int32_t>(NpuArch::DAV_2201)}), 95);
4539-REGISTER_TILING_TEMPLATE_WITH_SOCVERSION(FlashAttentionScore, FlashAttentionScoreTilingS1s2Bn2gs1, std::vector<int32_t>({static_cast<int32_t>(platform_ascendc::SocVersion::ASCEND910B), static_cast<int32_t>(platform_ascendc::SocVersion::ASCEND910_93)}), 96);4539+REGISTER_TILING_TEMPLATE_WITH_ARCH(FlashAttentionScore, FlashAttentionScoreTilingS1s2Bn2gs1, std::vector<int32_t>({static_cast<int32_t>(NpuArch::DAV_2201)}), 96);
4540-REGISTER_TILING_TEMPLATE_WITH_SOCVERSION(FlashAttentionScore, FlashAttentionScoreTilingS1Bn2gs1, std::vector<int32_t>({static_cast<int32_t>(platform_ascendc::SocVersion::ASCEND910B), static_cast<int32_t>(platform_ascendc::SocVersion::ASCEND910_93)}), 97);4540+REGISTER_TILING_TEMPLATE_WITH_ARCH(FlashAttentionScore, FlashAttentionScoreTilingS1Bn2gs1, std::vector<int32_t>({static_cast<int32_t>(NpuArch::DAV_2201)}), 97);
4541-REGISTER_TILING_TEMPLATE_WITH_SOCVERSION(FlashAttentionScore, FlashAttentionScoreTilingB, std::vector<int32_t>({static_cast<int32_t>(platform_ascendc::SocVersion::ASCEND910B), static_cast<int32_t>(platform_ascendc::SocVersion::ASCEND910_93)}), 98);4541+REGISTER_TILING_TEMPLATE_WITH_ARCH(FlashAttentionScore, FlashAttentionScoreTilingB, std::vector<int32_t>({static_cast<int32_t>(NpuArch::DAV_2201)}), 98);
4542} // namespace FA4542} // namespace FA
4543} // namespace optiling4543} // namespace optiling
Mattention/flash_attention_score/op_host/arch35/flash_attention_score_tiling_regbase.h+1-0
@@ -409,6 +409,7 @@ protected:
409 uint32_t aivNum;409 uint32_t aivNum;
410 uint32_t aicNum;410 uint32_t aicNum;
411 platform_ascendc::SocVersion socVersion;411 platform_ascendc::SocVersion socVersion;
412+ NpuArch npuArch = NpuArch::DAV_RESV;
412 413 
413 matmul_tiling::DataType bmmDtype = matmul_tiling::DataType::DT_FLOAT;414 matmul_tiling::DataType bmmDtype = matmul_tiling::DataType::DT_FLOAT;
414 matmul_tiling::DataType bmm1OutDtype = matmul_tiling::DataType::DT_FLOAT;415 matmul_tiling::DataType bmm1OutDtype = matmul_tiling::DataType::DT_FLOAT;
Mattention/flash_attention_score/op_host/flash_attention_score_tiling.cpp+5-4
@@ -355,19 +355,19 @@ ASCENDC_EXTERN_C ge::graphStatus TilingFlashAttentionScore(gert::TilingContext *
355 return ge::GRAPH_FAILED);355 return ge::GRAPH_FAILED);
356 356
357 auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr);357 auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr);
358- if (ascendcPlatform.GetSocVersion() == platform_ascendc::SocVersion::ASCEND910_95) {358+ if (ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_3510) {
359- OP_LOGW(context, "Current soc version is ASCEND910_95.");359+ OP_LOGW(context, "Current npu arch is dav-3510.");
360 if (IsEmptyInputRegbase(context)) {360 if (IsEmptyInputRegbase(context)) {
361 return ge::GRAPH_SUCCESS;361 return ge::GRAPH_SUCCESS;
362 }362 }
363 } else {363 } else {
364- OP_LOGW(context, "Current soc version is not ASCEND910_95.");364+ OP_LOGW(context, "Current npu arch is not dav-3510.");
365 if (IsEmptyInput(context)) {365 if (IsEmptyInput(context)) {
366 return ge::GRAPH_SUCCESS;366 return ge::GRAPH_SUCCESS;
367 }367 }
368 }368 }
369 369 
370- auto resultCode = TilingRegistryNew::GetInstance().DoTilingImpl(context);370+ auto resultCode = TilingRegistryArch::GetInstance().DoTilingImpl(context);
371 return resultCode;371 return resultCode;
372}372}
373 373 
@@ -386,6 +386,7 @@ ASCENDC_EXTERN_C ge::graphStatus TilingPrepareForFlashAttentionScore(gert::Tilin
386 compileInfoPtr->aivNum = ascendcPlatform.GetCoreNumAiv();386 compileInfoPtr->aivNum = ascendcPlatform.GetCoreNumAiv();
387 compileInfoPtr->aicNum = ascendcPlatform.GetCoreNumAic();387 compileInfoPtr->aicNum = ascendcPlatform.GetCoreNumAic();
388 compileInfoPtr->socVersion = ascendcPlatform.GetSocVersion();388 compileInfoPtr->socVersion = ascendcPlatform.GetSocVersion();
389+ compileInfoPtr->npuArch = ascendcPlatform.GetCurNpuArch();
389 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, compileInfoPtr->ubSize);390 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, compileInfoPtr->ubSize);
390 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L1, compileInfoPtr->l1Size);391 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L1, compileInfoPtr->l1Size);
391 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_C, compileInfoPtr->l0cSize);392 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_C, compileInfoPtr->l0cSize);
Mattention/flash_attention_score/op_host/flash_attention_score_tiling_common.h+1-0
@@ -28,6 +28,7 @@ struct FlashAttentionScoreCompileInfo {
28 uint64_t l0cSize;28 uint64_t l0cSize;
29 uint64_t l2CacheSize;29 uint64_t l2CacheSize;
30 platform_ascendc::SocVersion socVersion;30 platform_ascendc::SocVersion socVersion;
31+ NpuArch npuArch;
31};32};
32 33 
33} // namespace optiling34} // namespace optiling
Aattention/flash_attention_score/op_kernel/arch32/basic_modules/common_header.h+133-0
@@ -0,0 +1,133 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 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+/*!
12+ * \file common_header.h
13+ * \brief
14+ */
15+ 
16+#ifndef _COMMON_HEADER_H_
17+#define _COMMON_HEADER_H_
18+ 
19+#include "kernel_operator.h"
20+ 
21+#define SET_FLAG(trigger, waiter, e) AscendC::SetFlag<AscendC::HardEvent::trigger##_##waiter>((e))
22+#define WAIT_FLAG(trigger, waiter, e) AscendC::WaitFlag<AscendC::HardEvent::trigger##_##waiter>((e))
23+#define PIPE_BARRIER(pipe) AscendC::PipeBarrier<PIPE_##pipe>()
24+ 
25+constexpr int32_t BASE_BLOCK_LENGTH = 128;
26+constexpr int32_t C0_SIZE = 16;
27+constexpr int32_t SIZE_16 = 16;
28+constexpr int32_t SIZE_32 = 32;
29+constexpr int32_t SIZE_64 = 64;
30+constexpr int32_t SIZE_128 = 128;
31+constexpr int32_t SIZE_256 = 256;
32+constexpr int32_t SIZE_384 = 384;
33+constexpr int32_t SIZE_512 = 512;
34+constexpr int32_t SIZE_ONE_K = 1024;
35+constexpr int32_t SIZE_LONG_BLOCK = 16384;
36+constexpr int32_t BASE_M_128 = 128;
37+constexpr int32_t BASE_N_128 = 128;
38+ 
39+enum class ArchType { ASCEND_V220, ASCEND_V200, ASCEND_M200 };
40+enum class BufferType { ASCEND_UB, ASCEND_CB, ASCEND_L0A, ASCEND_L0B, ASCEND_L0C, ASCEND_MAX };
41+ 
42+template <ArchType ArchTag>
43+struct HardwareInfo {
44+ static uint32_t const l2BW = 5;
45+ static uint32_t const hbmBW = 1;
46+ static uint32_t const supportMix = 0;
47+ static uint32_t const l1Size = 512 * 1024;
48+ static uint32_t const l0ASize = 64 * 1024;
49+ static uint32_t const l0BSize = 64 * 1024;
50+ static uint32_t const l0CSize = 128 * 1024;
51+ static uint32_t const l2Size = 192 * 1024 * 1024;
52+ static uint32_t const biasSize = 1024;
53+ static uint32_t const fixBufSize = 7 * 1024;
54+ static uint32_t const ubSize = 192 * 1024;
55+ static uint32_t const fractalSize = 512;
56+ static uint32_t const l1l0BlockSize = 32;
57+ static uint32_t const btBlockSize = 64;
58+ static uint32_t const fbBlockSize = 128;
59+};
60+ 
61+ 
62+template <ArchType ArchTag>
63+struct AsdopsBuffer {
64+public:
65+ __aicore__ AsdopsBuffer()
66+ {
67+ constexpr uint32_t bufferSize[(uint32_t)BufferType::ASCEND_MAX] = {HardwareInfo<ArchTag>::ubSize,
68+ HardwareInfo<ArchTag>::l1Size,
69+ HardwareInfo<ArchTag>::l0ASize,
70+ HardwareInfo<ArchTag>::l0BSize,
71+ HardwareInfo<ArchTag>::l0CSize};
72+#ifdef __DAV_C220_VEC__
73+ tensor[(uint32_t)BufferType::ASCEND_UB].InitBuffer(0, bufferSize[(uint32_t)BufferType::ASCEND_UB]);
74+ tensor[(uint32_t)BufferType::ASCEND_UB].address_.logicPos = static_cast<uint8_t>(AscendC::TPosition::VECIN);
75+#elif __DAV_C220_CUBE__
76+ tensor[(uint32_t)BufferType::ASCEND_CB].InitBuffer(0, bufferSize[(uint32_t)BufferType::ASCEND_CB]);
77+ tensor[(uint32_t)BufferType::ASCEND_CB].address_.logicPos = static_cast<uint8_t>(AscendC::TPosition::A1);
78+ tensor[(uint32_t)BufferType::ASCEND_L0A].InitBuffer(0, bufferSize[(uint32_t)BufferType::ASCEND_L0A]);
79+ tensor[(uint32_t)BufferType::ASCEND_L0A].address_.logicPos = static_cast<uint8_t>(AscendC::TPosition::A2);
80+ tensor[(uint32_t)BufferType::ASCEND_L0B].InitBuffer(0, bufferSize[(uint32_t)BufferType::ASCEND_L0B]);
81+ tensor[(uint32_t)BufferType::ASCEND_L0B].address_.logicPos = static_cast<uint8_t>(AscendC::TPosition::B2);
82+ tensor[(uint32_t)BufferType::ASCEND_L0C].InitBuffer(0, bufferSize[(uint32_t)BufferType::ASCEND_L0C]);
83+ tensor[(uint32_t)BufferType::ASCEND_L0C].address_.logicPos = static_cast<uint8_t>(AscendC::TPosition::CO1);
84+#else
85+ tensor[(uint32_t)BufferType::ASCEND_UB].InitBuffer(0, bufferSize[(uint32_t)BufferType::ASCEND_UB]);
86+ tensor[(uint32_t)BufferType::ASCEND_UB].address_.logicPos = static_cast<uint8_t>(AscendC::TPosition::VECIN);
87+ tensor[(uint32_t)BufferType::ASCEND_CB].InitBuffer(0, bufferSize[(uint32_t)BufferType::ASCEND_CB]);
88+ tensor[(uint32_t)BufferType::ASCEND_CB].address_.logicPos = static_cast<uint8_t>(AscendC::TPosition::A1);
89+ tensor[(uint32_t)BufferType::ASCEND_L0A].InitBuffer(0, bufferSize[(uint32_t)BufferType::ASCEND_L0A]);
90+ tensor[(uint32_t)BufferType::ASCEND_L0A].address_.logicPos = static_cast<uint8_t>(AscendC::TPosition::A2);
91+ tensor[(uint32_t)BufferType::ASCEND_L0B].InitBuffer(0, bufferSize[(uint32_t)BufferType::ASCEND_L0B]);
92+ tensor[(uint32_t)BufferType::ASCEND_L0B].address_.logicPos = static_cast<uint8_t>(AscendC::TPosition::B2);
93+ tensor[(uint32_t)BufferType::ASCEND_L0C].InitBuffer(0, bufferSize[(uint32_t)BufferType::ASCEND_L0C]);
94+ tensor[(uint32_t)BufferType::ASCEND_L0C].address_.logicPos = static_cast<uint8_t>(AscendC::TPosition::CO1);
95+#endif
96+ };
97+ 
98+ template <BufferType BufferType_, typename DstDataType = half>
99+ __aicore__ AscendC::LocalTensor<DstDataType> GetBuffer(const uint32_t offset) const
100+ {
101+ return tensor[(uint32_t)BufferType_][offset].template ReinterpretCast<DstDataType>();
102+ }
103+ 
104+public:
105+ AscendC::LocalTensor<uint8_t> tensor[(uint32_t)BufferType::ASCEND_MAX];
106+};
107+ 
108+struct AddrInfo {
109+ uint64_t left;
110+ uint64_t right;
111+ uint64_t out;
112+ int32_t kx = 0;
113+ int32_t ky = 0;
114+ int32_t lineStride = 0;
115+ bool lowerLeft;
116+ bool upperRight;
117+ int32_t S1Idx;
118+ int32_t S2Idx;
119+ int32_t blockStart;
120+};
121+ 
122+ 
123+template <typename T>
124+inline __aicore__ T RoundUp(const T val, const T align)
125+{
126+ static_assert(std::is_arithmetic<T>::value, "T must be an arithmetic type");
127+ if (align == 0 || val + align - 1 < val) {
128+ return val;
129+ }
130+ return (val + align - 1) / align * align;
131+}
132+ 
133+#endif // COMMON_HEADER_H
Aattention/flash_attention_score/op_kernel/arch32/basic_modules/cube_modules/cube1_op.h+273-0
@@ -0,0 +1,273 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 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+/*!
12+ * \file cube1_op.h
13+ * \brief
14+ */
15+ 
16+ 
17+#ifndef _CUBE1_OP_H_
18+#define _CUBE1_OP_H_
19+ 
20+template <typename TYPE, LayOutTypeEnum layOutType>
21+__aicore__ inline void
22+CubeOp1<TYPE, layOutType>::Cube1Compute(const AddrInfo &shapeInfo, __gm__ TYPE* left, __gm__ TYPE* right, __gm__ float* out, bool needNz2Nd)
23+{
24+ uint64_t gm2L1SrcDValueA = 0;
25+ uint64_t gm2L1SrcDValueB = 0;
26+ if constexpr (layOutType == LayOutTypeEnum::LAYOUT_TND) {
27+ gm2L1SrcDValueA = headDim * qHeadNum;
28+ gm2L1SrcDValueB = headDim * kvHeadNum;
29+ } else if constexpr (layOutType == LayOutTypeEnum::LAYOUT_BSH) {
30+ gm2L1SrcDValueA = headDim * qHeadNum;
31+ gm2L1SrcDValueB = headDim * kvHeadNum;
32+ } else if constexpr (layOutType == LayOutTypeEnum::LAYOUT_SBH) {
33+ gm2L1SrcDValueA = bNum * headDim * qHeadNum;
34+ gm2L1SrcDValueB = bNum * headDim * kvHeadNum;
35+ } else if constexpr (layOutType == LayOutTypeEnum::LAYOUT_BNSD) {
36+ gm2L1SrcDValueA = headDim;
37+ gm2L1SrcDValueB = headDim;
38+ }
39+ int32_t km = shapeInfo.ky;
40+ int32_t kn = shapeInfo.kx;
41+ 
42+ int32_t l1_m_size_ = km;
43+ 
44+ int32_t n_loop = CeilDiv(kn, BASE_N_128);
45+ 
46+ auto gm_a = left + shapeInfo.left;
47+ auto gm_b = right + shapeInfo.right;
48+ auto gm_c = out;
49+ 
50+ uint64_t left_offset = gm_a - (__gm__ TYPE*)0;
51+ uint64_t right_offset = gm_b - (__gm__ TYPE*)0;
52+ uint64_t out_offset = gm_c - (__gm__ float*)0;
53+ 
54+ LocalTensor<TYPE>* l1_a_tensor = ping_pong_flag_l1_a_ ? &l1_a_pong_tensor : &l1_a_ping_tensor;
55+ 
56+ Cube1LoadDataAToL1(*l1_a_tensor, temp_tensor_bf16[left_offset], l1_m_size_, gm2L1SrcDValueA);
57+ 
58+ for (int n_index = 0; n_index < n_loop; n_index++) {
59+ int32_t l1_n_size_ = (n_index == n_loop - 1) ? (kn - n_index * BASE_N_128) : BASE_N_128;
60+ LocalTensor<TYPE>* l1_b_tensor = ping_pong_flag_l1_b_ ? &l1_b_pong_tensor : &l1_b_ping_tensor;
61+ 
62+ Cube1LoadDataBToL1(*l1_b_tensor, temp_tensor_bf16[right_offset + n_index * BASE_N_128 * gm2L1SrcDValueB],
63+ l1_n_size_, gm2L1SrcDValueB);
64+ 
65+ int32_t l1_m_size_align_ = RoundUp(l1_m_size_, C0_SIZE);
66+ int32_t l1_n_size_align_ = RoundUp(l1_n_size_, C0_SIZE);
67+ 
68+ for (int n_offset = 0; n_offset < l1_n_size_; n_offset += BASE_BLOCK_LENGTH) {
69+ int32_t n_mad_ = Min((l1_n_size_ - n_offset), BASE_BLOCK_LENGTH);
70+ int32_t n0_ = RoundUp(n_mad_, C0_SIZE);
71+ 
72+ LocalTensor<TYPE>* l0_b_tensor = ping_pong_flag_l0_b_ ? &l0_b_pong_tensor : &l0_b_ping_tensor;
73+ 
74+ Cube1LoadDataBToL0((*l0_b_tensor), (*l1_b_tensor), n0_, n_offset, l1_n_size_align_, l1_n_size_);
75+ 
76+ for (int m_offset = 0; m_offset < l1_m_size_; m_offset += SIZE_128) {
77+ int32_t m_mad_ = Min((l1_m_size_ - m_offset), BASE_BLOCK_LENGTH);
78+ int32_t m0_ = RoundUp(m_mad_, C0_SIZE);
79+ 
80+ LocalTensor<TYPE>* l0_a_tensor = ping_pong_flag_l0_a_ ? &l0_a_pong_tensor : &l0_a_ping_tensor;
81+ 
82+ Cube1LoadDataAToL0((*l0_a_tensor), (*l1_a_tensor), l1_m_size_align_, m0_, headDim, m_offset);
83+
84+ LocalTensor<float>* l0_c_tensor = ping_pong_flag_l0_c_ ? &l0_c_pong_tensor : &l0_c_ping_tensor;
85+ Cube1Mmad(*l0_c_tensor, *l0_a_tensor, *l0_b_tensor, m_mad_, n_mad_);
86+ 
87+ SET_FLAG(M, MTE1, FLAG_SHIFT + ping_pong_flag_l0_a_);
88+ 
89+ Cube1CopyOut(temp_tensor_fp32, *l0_c_tensor, out_offset, m_mad_, m0_,
90+ n0_, l1_m_size_, n_index, m_offset, n_mad_, kn, needNz2Nd);
91+ 
92+ ping_pong_flag_l0_c_ = 1 - ping_pong_flag_l0_c_;
93+ ping_pong_flag_l0_a_ = 1 - ping_pong_flag_l0_a_;
94+ }
95+ SET_FLAG(M, MTE1, FLAG_SHIFT + ping_pong_flag_l0_b_ + 2);
96+ ping_pong_flag_l0_b_ = 1 - ping_pong_flag_l0_b_;
97+ }
98+ SET_FLAG(MTE1, MTE2, ping_pong_flag_l1_b_);
99+ ping_pong_flag_l1_b_ = 1 - ping_pong_flag_l1_b_;
100+ }
101+ SET_FLAG(MTE1, MTE2, ping_pong_flag_l1_a_ + 2);
102+ ping_pong_flag_l1_a_ = 1 - ping_pong_flag_l1_a_;
103+}
104+ 
105+template <typename TYPE, LayOutTypeEnum layOutType>
106+__aicore__ inline void
107+CubeOp1<TYPE, layOutType>::Cube1LoadDataBToL1(LocalTensor<TYPE> dstTensor,
108+ GlobalTensor<TYPE> srcTensor,
109+ int32_t l1_n_size_, uint64_t gm2L1SrcDValueB)
110+{
111+ int32_t l1_n_size_align_ = RoundUp(l1_n_size_, C0_SIZE);
112+ commonNd2NzParams.nValue = l1_n_size_;
113+ commonNd2NzParams.dValue = headDim;
114+ commonNd2NzParams.srcDValue = gm2L1SrcDValueB;
115+ commonNd2NzParams.dstNzC0Stride = l1_n_size_align_;
116+ WAIT_FLAG(MTE1, MTE2, ping_pong_flag_l1_b_);
117+ AscendC::DataCopy(
118+ dstTensor,
119+ srcTensor,
120+ commonNd2NzParams
121+ );
122+ SET_FLAG(MTE2, MTE1, ping_pong_flag_l1_b_);
123+ WAIT_FLAG(MTE2, MTE1, ping_pong_flag_l1_b_);
124+}
125+ 
126+template <typename TYPE, LayOutTypeEnum layOutType>
127+__aicore__ inline void
128+CubeOp1<TYPE, layOutType>::Cube1LoadDataAToL1(LocalTensor<TYPE> dstTensor,
129+ GlobalTensor<TYPE> srcTensor,
130+ int32_t l1_m_size_, uint64_t gm2L1SrcDValueA)
131+{
132+ int32_t l1_m_size_align_ = RoundUp(l1_m_size_, C0_SIZE);
133+ commonNd2NzParams.nValue = l1_m_size_;
134+ commonNd2NzParams.dValue = headDim;
135+ commonNd2NzParams.srcDValue = gm2L1SrcDValueA;
136+ commonNd2NzParams.dstNzC0Stride = l1_m_size_align_;
137+ WAIT_FLAG(MTE1, MTE2, ping_pong_flag_l1_a_ + 2);
138+ AscendC::DataCopy(
139+ dstTensor,
140+ srcTensor,
141+ commonNd2NzParams
142+ );
143+ SET_FLAG(MTE2, MTE1, ping_pong_flag_l1_a_ + 2);
144+ WAIT_FLAG(MTE2, MTE1, ping_pong_flag_l1_a_ + 2);
145+}
146+ 
147+template <typename TYPE, LayOutTypeEnum layOutType>
148+__aicore__ inline void CubeOp1<TYPE, layOutType>::Cube1CopyOut(GlobalTensor<float> dstTensor,
149+ LocalTensor<float> srcTensor, uint64_t gm_out_offset, int32_t m_mad_,
150+ int32_t m0_, int32_t n0_, int32_t l1_m_size_,
151+ int32_t n_index, int m_offset, int32_t n_mad_, int32_t kn,
152+ bool needNz2Nd)
153+{
154+ commonFixpipeParamsV220.mSize = m_mad_;
155+ commonFixpipeParamsV220.srcStride = m0_;
156+ 
157+ if (unlikely(needNz2Nd)) {
158+ // NZ出
159+ commonFixpipeParamsV220.nSize = n0_;
160+ commonFixpipeParamsV220.dstStride = l1_m_size_ * 2;
161+ auto out_offset = n_index * l1_m_size_ * BASE_N_128 + m_offset * SIZE_16;
162+ AscendC::Fixpipe<float, float, AscendC::CFG_NZ>(dstTensor[gm_out_offset + out_offset],
163+ srcTensor, commonFixpipeParamsV220);
164+ } else {
165+ // ND出
166+ commonFixpipeParamsV220.nSize = n_mad_;
167+ commonFixpipeParamsV220.dstStride = kn;
168+ auto out_offset = kn * m_offset + n_index * BASE_N_128;
169+ AscendC::Fixpipe<float, float, AscendC::CFG_ROW_MAJOR>(dstTensor[gm_out_offset + out_offset],
170+ srcTensor, commonFixpipeParamsV220);
171+ }
172+}
173+ 
174+template <typename TYPE, LayOutTypeEnum layOutType>
175+__aicore__ inline void CubeOp1<TYPE, layOutType>::Cube1Mmad(LocalTensor<float> dstCTensor,
176+ LocalTensor<TYPE> srcATensor,
177+ LocalTensor<TYPE> srcBTensor,
178+ int32_t m_mad_, int32_t n_mad_)
179+{
180+ uint16_t m_modify = (m_mad_ == 1) ? 2 : m_mad_;
181+ commonMadParams.m = m_modify;
182+ commonMadParams.n = n_mad_;
183+ commonMadParams.k = headDim;
184+ commonMadParams.unitFlag = 3;
185+ commonMadParams.cmatrixInitVal = true;
186+ AscendC::Mmad(
187+ dstCTensor,
188+ srcATensor,
189+ srcBTensor,
190+ commonMadParams
191+ );
192+}
193+ 
194+template <typename TYPE, LayOutTypeEnum layOutType>
195+__aicore__ inline void CubeOp1<TYPE, layOutType>::Cube1LoadDataAToL0(LocalTensor<TYPE> dstTensor,
196+ LocalTensor<TYPE> srcTensor,
197+ int32_t l1_m_size_align_, int32_t m0_, uint64_t headDim, int m_offset)
198+{
199+ WAIT_FLAG(M, MTE1, FLAG_SHIFT + ping_pong_flag_l0_a_);
200+ commonLoadData2dParamsNoTranspose.repeatTimes = headDim / SIZE_16;
201+ commonLoadData2dParamsNoTranspose.srcStride = l1_m_size_align_ / SIZE_16;
202+ for (int32_t i = 0; i < m0_ / SIZE_16; i++) {
203+ AscendC::LoadData(
204+ dstTensor[i * headDim * SIZE_16],
205+ srcTensor[m_offset * SIZE_16 + i * SIZE_256],
206+ commonLoadData2dParamsNoTranspose
207+ );
208+ }
209+ SET_FLAG(MTE1, M, ping_pong_flag_l0_a_);
210+ WAIT_FLAG(MTE1, M, ping_pong_flag_l0_a_);
211+}
212+ 
213+template <typename TYPE, LayOutTypeEnum layOutType>
214+__aicore__ inline void CubeOp1<TYPE, layOutType>::Cube1LoadDataBToL0(LocalTensor<TYPE> dstTensor,
215+ LocalTensor<TYPE> srcTensor,
216+ int32_t n0_, int n_offset, int32_t l1_n_size_align_, int32_t l1_n_size_)
217+{
218+ WAIT_FLAG(M, MTE1, FLAG_SHIFT + ping_pong_flag_l0_b_ + 2);
219+ if (l1_n_size_ == BASE_BLOCK_LENGTH) {
220+ commonLoadData2dParamsNoTranspose.repeatTimes = headDim * n0_ / SIZE_256;
221+ commonLoadData2dParamsNoTranspose.srcStride = 1;
222+ AscendC::LoadData(
223+ dstTensor,
224+ srcTensor,
225+ commonLoadData2dParamsNoTranspose
226+ );
227+ } else {
228+ commonLoadData2dParamsNoTranspose.repeatTimes = n0_ / SIZE_16;
229+ commonLoadData2dParamsNoTranspose.srcStride = 1;
230+ for (int i = 0; i < headDim / SIZE_16; i++) {
231+ AscendC::LoadData(
232+ dstTensor[i * n0_ * SIZE_16],
233+ srcTensor[i * l1_n_size_align_ * SIZE_16 + n_offset * SIZE_16],
234+ commonLoadData2dParamsNoTranspose
235+ );
236+ }
237+ }
238+ SET_FLAG(MTE1, M, ping_pong_flag_l0_b_ + 2);
239+ WAIT_FLAG(MTE1, M, ping_pong_flag_l0_b_ + 2);
240+}
241+ 
242+template <typename TYPE, LayOutTypeEnum layOutType>
243+__aicore__ inline void CubeOp1<TYPE, layOutType>::SetFlag()
244+{
245+ SET_FLAG(MTE1, MTE2, EVENT_ID0);
246+ SET_FLAG(MTE1, MTE2, EVENT_ID1);
247+ SET_FLAG(MTE1, MTE2, EVENT_ID2);
248+ SET_FLAG(MTE1, MTE2, EVENT_ID3);
249+ SET_FLAG(MTE1, MTE2, EVENT_ID4);
250+ SET_FLAG(MTE1, MTE2, EVENT_ID5);
251+ 
252+ SET_FLAG(M, MTE1, EVENT_ID3);
253+ SET_FLAG(M, MTE1, EVENT_ID4);
254+ SET_FLAG(M, MTE1, EVENT_ID5);
255+ SET_FLAG(M, MTE1, EVENT_ID6);
256+}
257+ 
258+template <typename TYPE, LayOutTypeEnum layOutType>
259+__aicore__ inline void CubeOp1<TYPE, layOutType>::WaitFlag()
260+{
261+ WAIT_FLAG(MTE1, MTE2, EVENT_ID0);
262+ WAIT_FLAG(MTE1, MTE2, EVENT_ID1);
263+ WAIT_FLAG(MTE1, MTE2, EVENT_ID2);
264+ WAIT_FLAG(MTE1, MTE2, EVENT_ID3);
265+ WAIT_FLAG(MTE1, MTE2, EVENT_ID4);
266+ WAIT_FLAG(MTE1, MTE2, EVENT_ID5);
267+ 
268+ WAIT_FLAG(M, MTE1, EVENT_ID3);
269+ WAIT_FLAG(M, MTE1, EVENT_ID4);
270+ WAIT_FLAG(M, MTE1, EVENT_ID5);
271+ WAIT_FLAG(M, MTE1, EVENT_ID6);
272+}
273+#endif
Aattention/flash_attention_score/op_kernel/arch32/basic_modules/cube_modules/cube2_op.h+238-0文件内容审核中，请稍后刷新重试
Aattention/flash_attention_score/op_kernel/arch32/basic_modules/cube_op1.h+246-0文件内容审核中，请稍后刷新重试
Mattention/flash_attention_score/op_kernel/arch32/flash_attention_score_s1s2_bn2gs1_sab.h+39-11
@@ -23,6 +23,8 @@
23#include "kernel_tiling/kernel_tiling.h"23#include "kernel_tiling/kernel_tiling.h"
24#include "lib/matmul_intf.h"24#include "lib/matmul_intf.h"
25#include "pse.h"25#include "pse.h"
26+#include "basic_modules/common_header.h"
27+#include "basic_modules/cube_op1.h"
26 28 
27using matmul::MatmulType;29using matmul::MatmulType;
28using AscendC::CrossCoreSetFlag;30using AscendC::CrossCoreSetFlag;
@@ -397,6 +399,9 @@ protected:
397 uint64_t SYNC_C1_V1_FLAG[3] = {4, 5, 6};399 uint64_t SYNC_C1_V1_FLAG[3] = {4, 5, 6};
398 uint64_t SYNC_V1_C2_FLAG[3] = {7, 8, 9};400 uint64_t SYNC_V1_C2_FLAG[3] = {7, 8, 9};
399 uint64_t SYNC_C2_V2_FLAG[3] = {1, 2, 3};401 uint64_t SYNC_C2_V2_FLAG[3] = {1, 2, 3};
402+ 
403+ // 基础API实现MATMUL
404+ CUBE_OP1::CubeOp1<INPUT_T, layOutType> cubeOp;
400};405};
401 406 
402template <ImplModeEnum implMode, LayOutTypeEnum layOutType, bool hasPse, bool hasAtten, bool hasDrop, typename INPUT_T,407template <ImplModeEnum implMode, LayOutTypeEnum layOutType, bool hasPse, bool hasAtten, bool hasDrop, typename INPUT_T,
@@ -597,6 +602,14 @@ __aicore__ inline void FlashAttentionScoreS1s2Bn2gs1SameAB<implMode, layOutType,
597 this->pipe->InitBuffer(this->softmaxTempBuf, vecS1BaseSize * blockBytes); // 16k602 this->pipe->InitBuffer(this->softmaxTempBuf, vecS1BaseSize * blockBytes); // 16k
598 }603 }
599 this->pipe->InitBuffer(this->stage1PongBuf, stage1PongSize); // i.a 34k604 this->pipe->InitBuffer(this->stage1PongBuf, stage1PongSize); // i.a 34k
605+ if constexpr (mmPolicyType != MmPolicyType::UNSPLITK) {
606+ if (this->dSize == 64) {
607+ cubeOp.Init(this->n2G,
608+ this->tilingData->inputParams.n2Size,
609+ this->dSize,
610+ this->tilingData->inputParams.bSize);
611+ }
612+ }
600}613}
601 614 
602template <ImplModeEnum implMode, LayOutTypeEnum layOutType, bool hasPse, bool hasAtten, bool hasDrop, typename INPUT_T,615template <ImplModeEnum implMode, LayOutTypeEnum layOutType, bool hasPse, bool hasAtten, bool hasDrop, typename INPUT_T,
@@ -641,7 +654,7 @@ __aicore__ inline void FlashAttentionScoreS1s2Bn2gs1SameAB<implMode, layOutType,
641 // 计算切分轴的乘积654 // 计算切分轴的乘积
642 this->s2BaseN2D = this->s2BaseSize * this->n2D;655 this->s2BaseN2D = this->s2BaseSize * this->n2D;
643 this->s2BaseNratioSize = this->s2BaseSize * this->tilingData->coreParams.nRatio;656 this->s2BaseNratioSize = this->s2BaseSize * this->tilingData->coreParams.nRatio;
644- 657+ 
645 if constexpr (hasRope == true) {658 if constexpr (hasRope == true) {
646 this->gDRope = this->tilingData->inputParams.gSize * dRopeSize;659 this->gDRope = this->tilingData->inputParams.gSize * dRopeSize;
647 this->n2DRope = this->tilingData->inputParams.n2Size * dRopeSize;660 this->n2DRope = this->tilingData->inputParams.n2Size * dRopeSize;
@@ -673,7 +686,7 @@ __aicore__ inline void FlashAttentionScoreS1s2Bn2gs1SameAB<implMode, layOutType,
673 this->mm1Ka2 = this->bN2GDRope;686 this->mm1Ka2 = this->bN2GDRope;
674 this->mm1Kb1 = this->bN2D;687 this->mm1Kb1 = this->bN2D;
675 this->mm1Kb2 = this->bN2DRope;688 this->mm1Kb2 = this->bN2DRope;
676- this->mm1Kb = this->mm1Kb1 + this->mm1Kb2;689+ this->mm1Kb = this->mm1Kb1 + this->mm1Kb2;
677 this->mm2Kb = this->bN2D2;690 this->mm2Kb = this->bN2D2;
678 } else if constexpr (layOutType == LayOutTypeEnum::LAYOUT_BNSD) {691 } else if constexpr (layOutType == LayOutTypeEnum::LAYOUT_BNSD) {
679 // BNSD692 // BNSD
@@ -682,7 +695,7 @@ __aicore__ inline void FlashAttentionScoreS1s2Bn2gs1SameAB<implMode, layOutType,
682 this->s2BaseNratioD = this->s2BaseNratioSize * this->dSize;695 this->s2BaseNratioD = this->s2BaseNratioSize * this->dSize;
683 this->mm1Ka1 = this->dSize;696 this->mm1Ka1 = this->dSize;
684 this->mm1Kb1 = this->dSize;697 this->mm1Kb1 = this->dSize;
685- this->mm1Kb = this->mm1Kb1; 698+ this->mm1Kb = this->mm1Kb1;
686 this->mm2Kb = this->d2Size;699 this->mm2Kb = this->d2Size;
687 }700 }
688 701 
@@ -1150,9 +1163,17 @@ FlashAttentionScoreS1s2Bn2gs1SameAB<implMode, layOutType, hasPse, hasAtten, hasD
1150 this->Bmm1SetTensorA(extraInfo, bmm1);1163 this->Bmm1SetTensorA(extraInfo, bmm1);
1151 this->SetBmm1TensorB(extraInfo, bmm1);1164 this->SetBmm1TensorB(extraInfo, bmm1);
1152 if constexpr (mmPolicyType != MmPolicyType::UNSPLITK) {1165 if constexpr (mmPolicyType != MmPolicyType::UNSPLITK) {
1153- bmm1.template IterateAll<false>(this->mm1Res[extraInfo.taskIdMod2], hasRope, false, true);1166+ if (this->dSize != 64)
1154- bmm1.End();1167+ {
1155- return;1168+ bmm1.template IterateAll<false>(this->mm1Res[extraInfo.taskIdMod2], hasRope, false, true);
1169+ bmm1.End();
1170+ return;
1171+ } else
1172+ {
1173+ cubeOp.Cube1Process(extraInfo.qCoreOffset, extraInfo.kCoreOffset, extraInfo.cubeS1RealSize, extraInfo.s2RealSize,
1174+ this->queryGm, this->keyGm, this->mm1Res[extraInfo.taskIdMod2], extraInfo.needNz2Nd);
1175+ return;
1176+ }
1156 }1177 }
1157 1178 
1158 int32_t aRowNum = extraInfo.cubeS1RealSize / mm1BaseM;1179 int32_t aRowNum = extraInfo.cubeS1RealSize / mm1BaseM;
@@ -2284,12 +2305,19 @@ __aicore__ inline void FlashAttentionScoreS1s2Bn2gs1SameAB<implMode, layOutType,
2284 bmm2.SetTensorA(this->stage1Res[extraInfo.taskIdMod2]);2305 bmm2.SetTensorA(this->stage1Res[extraInfo.taskIdMod2]);
2285 bmm2.SetTensorB(this->valueGm[vCoreOffset]);2306 bmm2.SetTensorB(this->valueGm[vCoreOffset]);
2286 bmm2.SetTail(extraInfo.cubeS1RealSize, this->d2Size, extraInfo.s2RealSize);2307 bmm2.SetTail(extraInfo.cubeS1RealSize, this->d2Size, extraInfo.s2RealSize);
2287- 2308+ if (this->dSize != 64)
2288- bmm2.template IterateAll<false>(this->mm2Res[extraInfo.taskIdMod2], false, false, true);2309+ {
2289- bmm2.End();2310+ bmm2.template IterateAll<false>(this->mm2Res[extraInfo.taskIdMod2], false, false, true);
2290- return;2311+ bmm2.End();
2312+ return;
2313+ } else {
2314+ cubeOp.Cube2Process(vCoreOffset, extraInfo.cubeS1RealSize, extraInfo.s2RealSize,
2315+ this->stage1Res[extraInfo.taskIdMod2], this->valueGm[vCoreOffset],
2316+ this->mm2Res[extraInfo.taskIdMod2]);
2317+ return;
2318+ }
2291 }2319 }
2292- 2320+
2293 LocalTensor<INPUT_T> scmATensor;2321 LocalTensor<INPUT_T> scmATensor;
2294 LocalTensor<INPUT_T> scmBTensor;2322 LocalTensor<INPUT_T> scmBTensor;
2295 2323
Mattention/flash_attention_score/op_kernel/arch35/flash_attention_score_drop_mask_adapter_regbase.h+2-2
@@ -16,7 +16,7 @@
16#ifndef FLASH_ATTENTION_SCORE_DROP_MASK_ADAPTER_REGBASE_H16#ifndef FLASH_ATTENTION_SCORE_DROP_MASK_ADAPTER_REGBASE_H
17#define FLASH_ATTENTION_SCORE_DROP_MASK_ADAPTER_REGBASE_H17#define FLASH_ATTENTION_SCORE_DROP_MASK_ADAPTER_REGBASE_H
18 18 
19-#include "kernel_operator.h"19+#include "kernel_basic_intf.h"
20#include "kernel_tiling/kernel_tiling.h"20#include "kernel_tiling/kernel_tiling.h"
21 21 
22class FlashAttentionScoreDropMaskAdapterRegbase {22class FlashAttentionScoreDropMaskAdapterRegbase {
@@ -98,7 +98,7 @@ __aicore__ inline void FlashAttentionScoreDropMaskAdapterRegbase::Process()
98 98 
99 int32_t baseUbCalSize = tilingData->dropmaskParamsRegbase.baseUbCalSize;99 int32_t baseUbCalSize = tilingData->dropmaskParamsRegbase.baseUbCalSize;
100 AscendC::LocalTensor<half> dropMaskSelSrc = dropMaskSelSrcTBuf.template Get<half>();100 AscendC::LocalTensor<half> dropMaskSelSrc = dropMaskSelSrcTBuf.template Get<half>();
101- AscendC::Duplicate<half>(dropMaskSelSrc, 1.0, baseUbCalSize);101+ AscendC::Duplicate<half>(dropMaskSelSrc, static_cast<half>(1.0), baseUbCalSize);
102 102 
103 AscendC::BinaryRepeatParams binaryRepeatParams;103 AscendC::BinaryRepeatParams binaryRepeatParams;
104 binaryRepeatParams.src0BlkStride = 1;104 binaryRepeatParams.src0BlkStride = 1;
Mattention/flash_attention_score/op_kernel/arch35/flash_attention_score_empty_tensor_regbase.h+5-1
@@ -16,10 +16,14 @@
16#ifndef FLASH_ATTENTION_SCORE_REGBASE_EMPTY_TENSOR_H16#ifndef FLASH_ATTENTION_SCORE_REGBASE_EMPTY_TENSOR_H
17#define FLASH_ATTENTION_SCORE_REGBASE_EMPTY_TENSOR_H17#define FLASH_ATTENTION_SCORE_REGBASE_EMPTY_TENSOR_H
18 18 
19-#include "kernel_operator.h"19+#include "kernel_basic_intf.h"
20#include "kernel_tiling/kernel_tiling.h"20#include "kernel_tiling/kernel_tiling.h"
21#include "flash_attention_score_template_tiling_key.h"21#include "flash_attention_score_template_tiling_key.h"
22+#if __has_include("../../../common/op_kernel/arch35/flash_attention_score_tiling_regbase.h")
22#include "../../../common/op_kernel/arch35/flash_attention_score_tiling_regbase.h"23#include "../../../common/op_kernel/arch35/flash_attention_score_tiling_regbase.h"
24+#else
25+#include "../../common/arch35/flash_attention_score_tiling_regbase.h"
26+#endif
23using namespace optiling;27using namespace optiling;
24 28 
25template <typename INPUT_T> class FlashAttentionScoreEmptyTensorRegbase {29template <typename INPUT_T> class FlashAttentionScoreEmptyTensorRegbase {
Mattention/flash_attention_score/op_kernel/arch35/flash_attention_score_entry_regbase.h+23-0
@@ -40,6 +40,28 @@
40 GET_TILING_DATA_WITH_STRUCT(FlashAttentionScoreSimplifiedTilingData, tilingDataIn, tiling); \40 GET_TILING_DATA_WITH_STRUCT(FlashAttentionScoreSimplifiedTilingData, tilingDataIn, tiling); \
41 const FlashAttentionScoreSimplifiedTilingData *__restrict tilingData = &tilingDataIn; \41 const FlashAttentionScoreSimplifiedTilingData *__restrict tilingData = &tilingDataIn; \
42 42 
43+#ifdef __CCE_KT_TEST__
44+#define INVOKE_FA_OP_IMPL_BASEAPI(templateClass, ...) \
45+ do { \
46+ __gm__ uint8_t *user = GetUserWorkspace(workspace); \
47+ REGBASE_COPY_TILING_DATA(tiling); \
48+ TPipe tPipe; \
49+ if (tilingData->inputParamsRegbase.needDropMaskOp) { \
50+ FlashAttentionScoreDropMaskAdapterRegbase dropMaskAdapter; \
51+ dropMaskAdapter.Init(dropMask, user, tilingData, &tPipe); \
52+ dropMaskAdapter.Process(); \
53+ tPipe.Reset(); \
54+ } \
55+ using CubeBlockType = typename BaseApi::FABlockCube<__VA_ARGS__>; \
56+ using VecBlockType = typename BaseApi::FABlockVecTrain<__VA_ARGS__>; \
57+ templateClass<CubeBlockType, VecBlockType> op; \
58+ op.InitBaseAPI(query, key, value, pse, dropMask, paddingMask, attenMask, prefix, actualSeqLengths, \
59+ actualSeqLengthsKv, nullptr, nullptr, nullptr, deqScaleQ, deqScaleK, deqScaleV, nullptr, \
60+ nullptr,queryRope, keyRope, softmaxMax, softmaxSum, softmaxOut, nullptr, attentionOut, user, \
61+ tilingData, &tPipe); \
62+ op.Process(); \
63+ } while (0)
64+#else
43#define INVOKE_FA_OP_IMPL_BASEAPI(templateClass, ...) \65#define INVOKE_FA_OP_IMPL_BASEAPI(templateClass, ...) \
44 do { \66 do { \
45 __gm__ uint8_t *user = GetUserWorkspace(workspace); \67 __gm__ uint8_t *user = GetUserWorkspace(workspace); \
@@ -61,6 +83,7 @@
61 op.Process(); \83 op.Process(); \
62 } while (0)84 } while (0)
63#endif85#endif
86+#endif
64 87 
65template<uint8_t implMode, uint8_t layout, uint16_t s1TemplateType, uint16_t s2TemplateType,88template<uint8_t implMode, uint8_t layout, uint16_t s1TemplateType, uint16_t s2TemplateType,
66 uint16_t dTemplateType, uint16_t dvTemplateType, uint8_t pseMode, bool hasAtten, bool hasDrop, bool hasRope,89 uint16_t dTemplateType, uint16_t dvTemplateType, uint8_t pseMode, bool hasAtten, bool hasDrop, bool hasRope,
Mattention/flash_attention_score/op_kernel/arch35/flash_attention_score_kernel_train.h+94-27
@@ -15,8 +15,13 @@
15 15 
16#ifndef FLASH_ATTENTION_SCORE_KERNEL_TRAIN_H_16#ifndef FLASH_ATTENTION_SCORE_KERNEL_TRAIN_H_
17#define FLASH_ATTENTION_SCORE_KERNEL_TRAIN_H_17#define FLASH_ATTENTION_SCORE_KERNEL_TRAIN_H_
18+#if __has_include("../../../common/op_kernel/arch35/flash_attention_score_kernel_base.h")
18#include "../../../common/op_kernel/arch35/flash_attention_score_kernel_base.h"19#include "../../../common/op_kernel/arch35/flash_attention_score_kernel_base.h"
19#include "../../../common/op_kernel/arch35/dropmask.h"20#include "../../../common/op_kernel/arch35/dropmask.h"
21+#else
22+#include "../../common/arch35/flash_attention_score_kernel_base.h"
23+#include "../../common/arch35/dropmask.h"
24+#endif
20namespace BaseApi {25namespace BaseApi {
21template <typename CubeBlockType, typename VecBlockType>26template <typename CubeBlockType, typename VecBlockType>
22class FlashAttentionScoreKernelTrain27class FlashAttentionScoreKernelTrain
@@ -35,6 +40,8 @@ private:
35 __aicore__ inline int64_t CalcRealTimes(int64_t relativePos, int64_t length);40 __aicore__ inline int64_t CalcRealTimes(int64_t relativePos, int64_t length);
36 __aicore__ inline int64_t CalcRealCoreIdx(int64_t relativePos, int64_t times, int64_t offsetCoreIdx,41 __aicore__ inline int64_t CalcRealCoreIdx(int64_t relativePos, int64_t times, int64_t offsetCoreIdx,
37 bool isPartialCalc);42 bool isPartialCalc);
43+ __aicore__ inline int64_t CalcRealCoreIdxVarlen(int64_t calcLoops, int64_t calcLoopsRemain,
44+ int64_t cycleCoreNums);
38};45};
39 46 
40template <typename CubeBlockType, typename VecBlockType>47template <typename CubeBlockType, typename VecBlockType>
@@ -95,6 +102,20 @@ __aicore__ inline int64_t FlashAttentionScoreKernelTrain<CubeBlockType, VecBlock
95 return realCoreIdx;102 return realCoreIdx;
96}103}
97 104 
105+template <typename CubeBlockType, typename VecBlockType>
106+__aicore__ inline int64_t FlashAttentionScoreKernelTrain<CubeBlockType, VecBlockType>::CalcRealCoreIdxVarlen(
107+ int64_t calcLoops, int64_t calcLoopsRemain, int64_t cycleCoreNums)
108+{
109+ int64_t realCoreIdx = 0;
110+ if (calcLoopsRemain == 0) {
111+ realCoreIdx = calcLoops * cycleCoreNums + this->aicIdx;
112+ } else {
113+ realCoreIdx = calcLoops * cycleCoreNums + (cycleCoreNums - this->aicIdx - 1);
114+ }
115+ 
116+ return realCoreIdx;
117+}
118+ 
98template <typename CubeBlockType, typename VecBlockType>119template <typename CubeBlockType, typename VecBlockType>
99__aicore__ inline int64_t FlashAttentionScoreKernelTrain<CubeBlockType, VecBlockType>::CalcRealTimes(120__aicore__ inline int64_t FlashAttentionScoreKernelTrain<CubeBlockType, VecBlockType>::CalcRealTimes(
100 int64_t relativePos, int64_t length)121 int64_t relativePos, int64_t length)
@@ -122,6 +143,17 @@ __aicore__ inline int64_t FlashAttentionScoreKernelTrain<CubeBlockType, VecBlock
122 * s1基本块4 -- core3143 * s1基本块4 -- core3
123 * s1基本块5 -- core2144 * s1基本块5 -- core2
124 * s1基本块6 -- core1145 * s1基本块6 -- core1
146+ *
147+ * 三、正倒序循环分核：将每个S1的基本块，先正序依次分发给各个核计算，再倒序依次分发给各个核计算，以此循环
148+ * s1基本块1 -- core1
149+ * s1基本块2 -- core2
150+ * s1基本块3 -- core3
151+ * s1基本块4 -- core3
152+ * s1基本块5 -- core2
153+ * s1基本块6 -- core1
154+ * s1基本块7 -- core1
155+ * s1基本块8 -- core2
156+ * ...
125 */157 */
126template <typename CubeBlockType, typename VecBlockType>158template <typename CubeBlockType, typename VecBlockType>
127__aicore__ inline void FlashAttentionScoreKernelTrain<CubeBlockType, VecBlockType>::Process()159__aicore__ inline void FlashAttentionScoreKernelTrain<CubeBlockType, VecBlockType>::Process()
@@ -130,8 +162,9 @@ __aicore__ inline void FlashAttentionScoreKernelTrain<CubeBlockType, VecBlockTyp
130 int64_t multiCoreInnerLimit = 0;162 int64_t multiCoreInnerLimit = 0;
131 163 
132 // 新分核模式164 // 新分核模式
133- // 1、S2全量计算的部分，采用顺序分核；165+ // 1、非TND场景 S2全量计算的部分，采用顺序分核；
134- // 2、S2部分计算的部分，采用对称分核：将N分成一半，上半部分顺序分核，下半部分与上半部分对称分核；166+ // 2、非TND场景 S2部分计算的部分，采用对称分核：将N分成一半，上半部分顺序分核，下半部分与上半部分对称分核；
167+ // 3、TND场景 正倒序循环分核
135 int64_t halfN = 0;168 int64_t halfN = 0;
136 int64_t partialCalcForwardNum = 0; // 当前核 在顺序部分计算中分配的S1方向上基本块个数；169 int64_t partialCalcForwardNum = 0; // 当前核 在顺序部分计算中分配的S1方向上基本块个数；
137 int64_t partialCalcReverseNum = 0; // 当前核 在倒序部分计算中分配的S1方向上基本块个数；170 int64_t partialCalcReverseNum = 0; // 当前核 在倒序部分计算中分配的S1方向上基本块个数；
@@ -140,26 +173,46 @@ __aicore__ inline void FlashAttentionScoreKernelTrain<CubeBlockType, VecBlockTyp
140 int64_t halfNCoreIdx = 0; // 下半部分第一个S1方向基本块对应的核索引；173 int64_t halfNCoreIdx = 0; // 下半部分第一个S1方向基本块对应的核索引；
141 int64_t partialCalcLength = this->sharedParams.firstFullLoadS1OuterIdx + 1; // 部分计算在单个S1上的长度；174 int64_t partialCalcLength = this->sharedParams.firstFullLoadS1OuterIdx + 1; // 部分计算在单个S1上的长度；
142 int64_t relativePosReverse = 0; // 当前核 与第一个S1方向基本块对应的核索引 相差的个数175 int64_t relativePosReverse = 0; // 当前核 与第一个S1方向基本块对应的核索引 相差的个数
176+ int64_t varlenCalcLoops = 0; // TND场景 需要进行计算的循环次数(正序+倒序为一次循环)
177+ int64_t varlenCalcLoopsRemain = 0;
178+ int64_t varlenCalcTimes = 0; // TND场景 需要计算的S1方向上基本块总数
179+ int64_t varlenCycleCoreNums = this->sharedParams.coreNum * 2; // TND场景 一次循环正序+倒序为两倍核数
180+ 
143 if (this->sharedParams.splitCoreMode == 1) {181 if (this->sharedParams.splitCoreMode == 1) {
144- const int64_t totalN = this->constInfo.n2G * this->sharedParams.bSize;182+ if (layout == LayOutTypeEnum::LAYOUT_TND) {
145- halfN = CeilDiv(totalN, 2);183+ varlenCalcLoops = this->sharedParams.totalSize / varlenCycleCoreNums;
146- int64_t partialCalcForwardLength = halfN * partialCalcLength;184+ varlenCalcLoopsRemain = this->sharedParams.totalSize % varlenCycleCoreNums;
147- int64_t partialCalcReverseLength = (totalN - halfN) * partialCalcLength;185+ varlenCalcTimes = varlenCalcLoops * 2;
148- int64_t fullCalcForwardLength = 0;186+ if (varlenCalcLoopsRemain >= this->aicIdx + 1) {
149- if (this->sharedParams.firstFullLoadS1OuterIdx == -1) {187+ varlenCalcTimes++;
150- fullCalcForwardLength = totalN * this->constInfo.s1OuterSize;188+ if (varlenCalcLoopsRemain > this->sharedParams.coreNum &&
189+ (this->aicIdx + 1) > varlenCycleCoreNums - varlenCalcLoopsRemain) {
190+ varlenCalcTimes++;
191+ }
192+ }
193+ multiCoreInnerOffset = 0;
194+ multiCoreInnerLimit = varlenCalcTimes;
151 } else {195 } else {
152- fullCalcForwardLength = totalN * (this->constInfo.s1OuterSize - partialCalcLength);196+ const int64_t totalN = this->constInfo.n2G * this->sharedParams.bSize;
197+ halfN = CeilDiv(totalN, 2);
198+ int64_t partialCalcForwardLength = halfN * partialCalcLength;
199+ int64_t partialCalcReverseLength = (totalN - halfN) * partialCalcLength;
200+ int64_t fullCalcForwardLength = 0;
201+ if (this->sharedParams.firstFullLoadS1OuterIdx == -1) {
202+ fullCalcForwardLength = totalN * this->constInfo.s1OuterSize;
203+ } else {
204+ fullCalcForwardLength = totalN * (this->constInfo.s1OuterSize - partialCalcLength);
205+ }
206+ halfNCoreIdx = (partialCalcForwardLength - 1) % this->sharedParams.coreNum;
207+ relativePosReverse = (halfNCoreIdx - this->aicIdx + this->sharedParams.coreNum) % this->sharedParams.coreNum;
208+ partialCalcForwardNum = CalcRealTimes(this->aicIdx, partialCalcForwardLength);
209+ partialCalcReverseNum = CalcRealTimes(relativePosReverse, partialCalcReverseLength);
210+ fullCalcForwardNum = CalcRealTimes(this->aicIdx, fullCalcForwardLength);
211+ partialCalcNum = partialCalcForwardNum + partialCalcReverseNum;
212+ // 表示当前核需要计算的次数
213+ multiCoreInnerOffset = 0;
214+ multiCoreInnerLimit = partialCalcForwardNum + partialCalcReverseNum + fullCalcForwardNum;
153 }215 }
154- halfNCoreIdx = (partialCalcForwardLength - 1) % this->sharedParams.coreNum;
155- relativePosReverse = (halfNCoreIdx - this->aicIdx + this->sharedParams.coreNum) % this->sharedParams.coreNum;
156- partialCalcForwardNum = CalcRealTimes(this->aicIdx, partialCalcForwardLength);
157- partialCalcReverseNum = CalcRealTimes(relativePosReverse, partialCalcReverseLength);
158- fullCalcForwardNum = CalcRealTimes(this->aicIdx, fullCalcForwardLength);
159- partialCalcNum = partialCalcForwardNum + partialCalcReverseNum;
160- // 表示当前核需要计算的次数
161- multiCoreInnerOffset = 0;
162- multiCoreInnerLimit = partialCalcForwardNum + partialCalcReverseNum + fullCalcForwardNum;
163 } else {216 } else {
164 // 表示核内切分起点217 // 表示核内切分起点
165 multiCoreInnerOffset = this->sharedParams.multiCoreInnerOffset;218 multiCoreInnerOffset = this->sharedParams.multiCoreInnerOffset;
@@ -169,7 +222,11 @@ __aicore__ inline void FlashAttentionScoreKernelTrain<CubeBlockType, VecBlockTyp
169 // 初始化AxisIdx222 // 初始化AxisIdx
170 RunParamStr<isInfer> runParam;223 RunParamStr<isInfer> runParam;
171 if constexpr (layout == LayOutTypeEnum::LAYOUT_TND) {224 if constexpr (layout == LayOutTypeEnum::LAYOUT_TND) {
172- CalS1OuterSize(multiCoreInnerOffset, runParam);225+ if (this->sharedParams.splitCoreMode == 1) {
226+ CalS1OuterSize(this->aicIdx, runParam);
227+ } else {
228+ CalS1OuterSize(multiCoreInnerOffset, runParam);
229+ }
173 }230 }
174 RunInfo<isInfer> runInfo[4];231 RunInfo<isInfer> runInfo[4];
175 int64_t taskId = 0;232 int64_t taskId = 0;
@@ -192,14 +249,20 @@ __aicore__ inline void FlashAttentionScoreKernelTrain<CubeBlockType, VecBlockTyp
192 } else {249 } else {
193 // 非最后三次伪循环，需要将当前核处理的次数转化为S1方向上基本块的索引值250 // 非最后三次伪循环，需要将当前核处理的次数转化为S1方向上基本块的索引值
194 if (this->sharedParams.splitCoreMode == 1) {251 if (this->sharedParams.splitCoreMode == 1) {
195- if (multiCoreInnerIdx >= 0 && multiCoreInnerIdx < partialCalcForwardNum) {252+ if (layout == LayOutTypeEnum::LAYOUT_TND) {
196- realCoreInnerIdx = CalcRealCoreIdx(this->aicIdx, multiCoreInnerIdx, 0, true);253+ int64_t curCalcLoops = multiCoreInnerIdx >> 1;
197- } else if (multiCoreInnerIdx >= partialCalcForwardNum && multiCoreInnerIdx < partialCalcNum) {254+ int64_t curCalcLoopsRemain = multiCoreInnerIdx & 1;
198- realCoreInnerIdx = CalcRealCoreIdx(relativePosReverse, multiCoreInnerIdx - partialCalcForwardNum,255+ realCoreInnerIdx = CalcRealCoreIdxVarlen(curCalcLoops, curCalcLoopsRemain, varlenCycleCoreNums);
199- halfN * this->constInfo.s1OuterSize, true);
200 } else {256 } else {
201- realCoreInnerIdx = CalcRealCoreIdx(this->aicIdx, multiCoreInnerIdx - partialCalcNum,257+ if (multiCoreInnerIdx >= 0 && multiCoreInnerIdx < partialCalcForwardNum) {
202- partialCalcLength, false);258+ realCoreInnerIdx = CalcRealCoreIdx(this->aicIdx, multiCoreInnerIdx, 0, true);
259+ } else if (multiCoreInnerIdx >= partialCalcForwardNum && multiCoreInnerIdx < partialCalcNum) {
260+ realCoreInnerIdx = CalcRealCoreIdx(relativePosReverse, multiCoreInnerIdx - partialCalcForwardNum,
261+ halfN * this->constInfo.s1OuterSize, true);
262+ } else {
263+ realCoreInnerIdx = CalcRealCoreIdx(this->aicIdx, multiCoreInnerIdx - partialCalcNum,
264+ partialCalcLength, false);
265+ }
203 }266 }
204 } else {267 } else {
205 realCoreInnerIdx = multiCoreInnerIdx;268 realCoreInnerIdx = multiCoreInnerIdx;
@@ -356,6 +419,10 @@ __aicore__ inline void FlashAttentionScoreKernelTrain<CubeBlockType, VecBlockTyp
356 if (this->sharedParams.sparseType == static_cast<uint8_t>(SparseModeEnum::CAUSAL)) {419 if (this->sharedParams.sparseType == static_cast<uint8_t>(SparseModeEnum::CAUSAL)) {
357 runParam.s2LineStartIdx = 0;420 runParam.s2LineStartIdx = 0;
358 runParam.s2LineEndIdx = Min((runParam.s1oIdx + 1) * this->s1BaseSize, actualS2Len);421 runParam.s2LineEndIdx = Min((runParam.s1oIdx + 1) * this->s1BaseSize, actualS2Len);
422+ } else if (this->sharedParams.sparseType == static_cast<uint8_t>(SparseModeEnum::RIGHT_DOWN_CAUSAL)) {
423+ runParam.s2LineStartIdx = 0;
424+ runParam.s2LineEndIdx =
425+ Min((runParam.s1oIdx + 1) * this->s1BaseSize + actualS2Len - actualS1Len, actualS2Len);
359 } else if (this->sharedParams.sparseType ==426 } else if (this->sharedParams.sparseType ==
360 static_cast<uint8_t>(SparseModeEnum::BAND)) {427 static_cast<uint8_t>(SparseModeEnum::BAND)) {
361 runParam.s2LineStartIdx = Max(428 runParam.s2LineStartIdx = Max(
Mattention/flash_attention_score/op_kernel/flash_attention_score_apt.cpp+2-2
@@ -13,11 +13,11 @@
13 * \brief13 * \brief
14 */14 */
15 15 
16-#include "kernel_operator.h"16+#include "kernel_basic_intf.h"
17#include "arch35/flash_attention_score_empty_tensor_regbase.h"17#include "arch35/flash_attention_score_empty_tensor_regbase.h"
18#include "arch35/flash_attention_score_template_tiling_key.h"18#include "arch35/flash_attention_score_template_tiling_key.h"
19#include "arch35/flash_attention_score_entry_regbase.h"19#include "arch35/flash_attention_score_entry_regbase.h"
20-#ifdef NOT_DYNAMIC_COMPILE20+#if __has_include("../../common/op_kernel/arch35/flash_attention_score_tiling_regbase.h")
21#include "../../common/op_kernel/arch35/flash_attention_score_tiling_regbase.h"21#include "../../common/op_kernel/arch35/flash_attention_score_tiling_regbase.h"
22#else22#else
23#include "../common/arch35/flash_attention_score_tiling_regbase.h"23#include "../common/arch35/flash_attention_score_tiling_regbase.h"
Aattention/flash_attention_score/tests/pytest/README.md+35-0
@@ -0,0 +1,35 @@
1+# FlashAttentionScore算子测试框架
2+ 
3+## 文件结构
4+ 
5+pytest/
6+- test_case.py # 测试用例集
7+- test_flash_attn.py # 执行主程序
8+- test_utils.py # 工具方法
9+- cpu_impl.py # cpu实现
10+- npu_impl.py # npu实现
11+ 
12+## 功能说明
13+ 
14+基于pytest测试框架，实现FA算子的功能验证：
15+- **CPU侧**：复现算子功能用以生成golden数据
16+- **NPU侧**：通过torch_npu进行算子直调获取实际数据
17+- **精度对比**：进行CPU与NPU结果的精度对比验证算子功能
18+ 
19+ 
20+## 环境配置
21+ 
22+### 前置要求
23+ 
24+1. 确认torch_npu为最新版本
25+2. source CANN包环境变量
26+ 
27+### Custom包调用
28+ 
29+支持custom包调用
30+ 
31+## 使用方法
32+ 
33+在pytest文件夹路径下执行：
34+ 
35+pytest -s
Aattention/flash_attention_score/tests/pytest/cpu_impl.py+235-0
@@ -0,0 +1,235 @@
1+#!/usr/bin/python
2+# -*- coding: utf-8 -*-
3+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+# See LICENSE in the root of the software repository for the full text of the License.
10+# ======================================================================================================================
11+ 
12+import torch
13+import torch_npu
14+import math
15+import numpy as np
16+import random
17+from test_utils import *
18+ 
19+ 
20+def tsoftmax(x):
21+ x_max = torch.max(x, dim=-1, keepdims=True)[0]
22+ x_sub = x.sub(x_max)
23+ y = torch.exp(x_sub)
24+ x_sum = y.sum(dim=-1, keepdims=True)
25+ res = y.div(x_sum)
26+ return res, x_max, x_sum
27+ 
28+ 
29+def skip_invalid_row(b, n, sq, softmax_res, x_max, x_sum):
30+ for i in range(b):
31+ for j in range(n):
32+ for k in range(sq):
33+ if x_max[i, j, k, :] == -40000.:
34+ softmax_res[i, j, k, :] = 0
35+ x_max[i, j, k, :] = torch.finfo(torch.float).min
36+ x_sum[i, j, k, :] = torch.finfo(torch.float).max
37+ return softmax_res, x_max, x_sum
38+ 
39+ 
40+def forward(q, k, v, drop_mask, atten_mask, pse_type, pse, scale, keep_prob, skip_invalid_row=False):
41+ q = q.float()
42+ k = k.float()
43+ v = v.float()
44+ 
45+ qk = torch.matmul(q, k.permute(0, 1, 3, 2))
46+ if pse is None:
47+ qk = qk.mul(scale)
48+ else:
49+ if pse_type == 1:
50+ qk = qk.add(pse).mul(scale)
51+ else:
52+ qk = qk.mul(scale).add(pse)
53+ if atten_mask is not None:
54+ qk = qk.masked_fill_(atten_mask.cpu().bool(), value=torch.tensor(-40000))
55+ softmax_res, x_max, x_sum = tsoftmax(qk)
56+ 
57+ if skip_invalid_row:
58+ b, n, sq, _ = softmax_res.shape
59+ softmax_res, x_max, x_sum = skip_invalid_row(b, n, sq, softmax_res, x_max, x_sum)
60+ 
61+ if drop_mask is not None:
62+ drop_res = softmax_res * drop_mask * (1.0 / keep_prob)
63+ else:
64+ drop_res = softmax_res
65+ 
66+ out = torch.matmul(drop_res, v)
67+ 
68+ x_max = x_max.broadcast_to(-1, -1, -1, 8)
69+ x_sum = x_sum.broadcast_to(-1, -1, -1, 8)
70+ return out, x_max, x_sum
71+ 
72+ 
73+def get_cu_seqlens(seqlens_list):
74+ cu = torch.zeros(len(seqlens_list) + 1, dtype=torch.int64)
75+ for i in range(1, len(seqlens_list) + 1):
76+ cu[i] = cu[i - 1] + seqlens_list[i - 1]
77+ return cu
78+ 
79+ 
80+def broadcastKV(n1, n2, kv_tensor, dtype):
81+ factor = n1 // n2
82+ kv_shape = kv_tensor.shape
83+ b = kv_shape[0]
84+ s = kv_shape[2]
85+ d = kv_shape[3]
86+ kv_res = torch.zeros(b, n1, s, d).to(dtype)
87+ for i in range(n1):
88+ j = i // factor
89+ kv_res[:, i:i + 1, :, :] = kv_tensor[:, j:j + 1, :, :]
90+ return kv_res
91+ 
92+ 
93+def tforward_normal(q, k, v, pse, **kwargs):
94+ b = kwargs.get("B")
95+ n1 = kwargs.get("N1")
96+ n2 = kwargs.get("N2", n1)
97+ s1 = kwargs.get("Sq")
98+ s2 = kwargs.get("Skv", s1)
99+ d = kwargs.get("D")
100+ sparse_mode = kwargs.get("sparse_mode", None)
101+ pre_tokens = kwargs.get("pre_tokens", 2147483647)
102+ next_tokens = kwargs.get("next_tokens", 2147483647)
103+ prefix = kwargs.get("prefix", [])
104+ seed = kwargs.get("seed", 0)
105+ offset = kwargs.get("offset", 0)
106+ 
107+ scale = kwargs.get("scale", 1 / (d ** 0.5))
108+ pse_type = kwargs.get("pse_type", 1)
109+ keep_prob = kwargs.get("keep_prob", 1)
110+ 
111+ if abs(1 - keep_prob) < 2e-14:
112+ drop_mask = None
113+ else:
114+ drop_mask = gen_dropmask(b, n1, s1, s2, keep_prob, seed, offset)
115+ 
116+ atten_mask = generate_cpu_mask(b, s1, s2, sparse_mode, pre_tokens, next_tokens, prefix)
117+ skip_invalid_row = False
118+ if atten_mask is not None:
119+ if sparse_mode == 0:
120+ skip_invalid_row = next_tokens < 0 or pre_tokens + s2 < s1
121+ elif sparse_mode == 3:
122+ skip_invalid_row = s1 > s2
123+ elif sparse_mode == 4:
124+ skip_invalid_row = pre_tokens < 0 or next_tokens + s2 < s1
125+ elif sparse_mode in [5, 6]:
126+ skip_invalid_row = True if 0 in prefix else False
127+ 
128+ if n1 != n2:
129+ k = broadcastKV(n1, n2, k, k.dtype)
130+ v = broadcastKV(n1, n2, v, v.dtype)
131+ out, x_max, x_sum = forward(q, k, v, drop_mask, atten_mask, pse_type, pse, scale, keep_prob, skip_invalid_row)
132+ return out, x_max, x_sum
133+ 
134+ 
135+def tforward_tnd(q, k, v, pse, **kwargs):
136+ actual_seq_qlen = list(kwargs.get("actual_seq_qlen"))
137+ actual_seq_kvlen = list(kwargs.get("actual_seq_kvlen", actual_seq_qlen))
138+ assert isinstance(actual_seq_qlen, list)
139+ assert isinstance(actual_seq_kvlen, list)
140+ seqlen_q_list = get_seqlen_list(actual_seq_qlen).numpy()
141+ seqlen_k_list = get_seqlen_list(actual_seq_kvlen).numpy()
142+ actual_seq_qlen.insert(0, 0)
143+ actual_seq_kvlen.insert(0, 0)
144+ input_layout = kwargs.get("input_layout")
145+ sparse_mode = kwargs.get("sparse_mode", None)
146+ pre_tokens = kwargs.get("pre_tokens", 2147483647)
147+ next_tokens = kwargs.get("next_tokens", 2147483647)
148+ prefix = kwargs.get("prefix", [])
149+ seed = kwargs.get("seed", 0)
150+ offset = kwargs.get("offset", 0)
151+ max_seqlen_q = seqlen_q_list.max()
152+ max_seqlen_k = seqlen_k_list.max()
153+ qk_size = seqlen_q_list * [math.ceil(i / 16) * 16 for i in seqlen_k_list]
154+ qk_pointer = get_cu_seqlens(qk_size).to(torch.int64)
155+ 
156+ b = len(seqlen_q_list)
157+ n1 = kwargs.get("N1")
158+ n2 = kwargs.get("N2", n1)
159+ d = kwargs.get("D")
160+ d_v = kwargs.get("DV", d)
161+ s1 = max(actual_seq_qlen)
162+ s2 = max(actual_seq_kvlen)
163+ 
164+ scale = kwargs.get("scale", 1 / (d ** 0.5))
165+ pse_type = kwargs.get("pse_type", 1)
166+ keep_prob = kwargs.get("keep_prob", 1)
167+ 
168+ pse_s1 = max(1024, max_seqlen_q)
169+ band_index = 0
170+ if sparse_mode == 7:
171+ for index in range(b - 1, -1, -1):
172+ if seqlen_q_list[index] != 0:
173+ band_index = index
174+ break
175+ elif sparse_mode == 8:
176+ for index in range(b):
177+ if seqlen_k_list[index] != 0:
178+ band_index = index
179+ break
180+ 
181+ if abs(1 - keep_prob) < 2e-14:
182+ drop_mask = None
183+ else:
184+ drop_mask = gen_dropmask_tnd(qk_pointer[-1].item() * n1, keep_prob, seed, offset)
185+ 
186+ out_golden = torch.zeros([1, n1, s1, d_v], dtype=q.dtype)
187+ x_max = torch.empty(0)
188+ x_sum = torch.empty(0)
189+ for i in range(b):
190+ if seqlen_q_list[i] != 0 and seqlen_q_list[i] != 0:
191+ qi = q[:, :, actual_seq_qlen[i]:actual_seq_qlen[i + 1]]
192+ ki = k[:, :, actual_seq_kvlen[i]:actual_seq_kvlen[i + 1]]
193+ vi = v[:, :, actual_seq_kvlen[i]:actual_seq_kvlen[i + 1]]
194+ 
195+ if n1 != n2:
196+ ki = broadcastKV(n1, n2, ki, ki.dtype)
197+ vi = broadcastKV(n1, n2, vi, vi.dtype)
198+ 
199+ if pse is None:
200+ psei = None
201+ else:
202+ pse_layout = kwargs.get("pse_layout").lower()
203+ assert pse_layout in ["bn", "n", "bnhs", "1nhs"]
204+ if pse_layout == "bn":
205+ psei = pse[i:i + 1, :, -seqlen_q_list[i]:, -seqlen_k_list[i]:]
206+ elif pse_layout == "n":
207+ psei = pse[:, :, -seqlen_q_list[i]:, -seqlen_k_list[i]:]
208+ elif pse_layout == "bnhs":
209+ psei = pse[i:i + 1, :, pse_s1 - seqlen_q_list[i]:pse_s1, max_seqlen_k - seqlen_k_list[i]:max_seqlen_k]
210+ elif pse_layout == "1nhs":
211+ psei = pse[:, :, pse_s1 - seqlen_q_list[i]:pse_s1, max_seqlen_k - seqlen_k_list[i]:max_seqlen_k]
212+ 
213+ if drop_mask is None:
214+ drop_maski = None
215+ else:
216+ drop_maski = drop_mask[(qk_pointer[i] * n1):(qk_pointer[i + 1] * n1)].reshape(n1, seqlen_q_list[i], math.ceil(seqlen_k_list[i] / 16) * 16)[:, :, :seqlen_k_list[i]]
217+
218+ if sparse_mode is None:
219+ atten_maski = None
220+ else:
221+ atten_maski = generate_cpu_mask(1, seqlen_q_list[i], seqlen_k_list[i], sparse_mode, pre_tokens, next_tokens, prefix, i, band_index)
222+ 
223+ outi, x_maxi, x_sumi = forward(qi, ki, vi, drop_maski, atten_maski, pse_type, psei, scale, keep_prob)
224+ out_golden[:, :, actual_seq_qlen[i]:actual_seq_qlen[i + 1]] = outi
225+ x_max = torch.cat((x_max, x_maxi.contiguous().view(-1)))
226+ x_sum = torch.cat((x_sum, x_sumi.contiguous().view(-1)))
227+ return out_golden, x_max, x_sum
228+ 
229+ 
230+def tforward(q, k, v, pse, **kwargs):
231+ input_layout=kwargs.get("input_layout")
232+ if input_layout == "TND":
233+ return tforward_tnd(q, k, v, pse, **kwargs)
234+ else:
235+ return tforward_normal(q, k, v, pse, **kwargs)
Aattention/flash_attention_score/tests/pytest/npu_impl.py+88-0
@@ -0,0 +1,88 @@
1+#!/usr/bin/python
2+# -*- coding: utf-8 -*-
3+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+# See LICENSE in the root of the software repository for the full text of the License.
10+# ======================================================================================================================
11+ 
12+import torch
13+import torch_npu
14+import math
15+import numpy as np
16+import random
17+from einops import rearrange
18+from test_utils import trans_bnsd_to_layout
19+ 
20+device_id = 0
21+device = torch.device(f"npu:{device_id}")
22+ 
23+ 
24+def fa_npu(q, k, v, q_rope, k_rope, atten_mask, pse, **kwargs):
25+ input_layout = kwargs.get("input_layout", "BNSD")
26+ n1 = kwargs.get("N1")
27+ d = kwargs.get("D")
28+ d_rope = kwargs.get("DRope", 0)
29+ sparse_mode = kwargs.get("sparse_mode", 0)
30+ prefix = kwargs.get("prefix", None)
31+ pre_tokens = kwargs.get("pre_tokens", 2147483647)
32+ next_tokens = kwargs.get("next_tokens", 2147483647)
33+ actual_seq_qlen = kwargs.get("actual_seq_qlen", None)
34+ actual_seq_kvlen = kwargs.get("actual_seq_kvlen", actual_seq_qlen)
35+ scale = kwargs.get("scale", 1 / (d ** 0.5))
36+ pse_type = kwargs.get("pse_type", 1)
37+ keep_prob = kwargs.get("keep_prob", 1)
38+ q_start_idx = kwargs.get("q_start_idx", None)
39+ kv_start_idx = kwargs.get("kv_start_idx", None)
40+ seed = kwargs.get("seed", 0)
41+ 
42+ q1 = trans_bnsd_to_layout(q, input_layout).contiguous().to(device)
43+ k1 = trans_bnsd_to_layout(k, input_layout).contiguous().to(device)
44+ v1 = trans_bnsd_to_layout(v, input_layout).contiguous().to(device)
45+ 
46+ if d_rope != 0:
47+ query_rope = trans_bnsd_to_layout(q_rope, input_layout).contiguous().to(device)
48+ key_rope = trans_bnsd_to_layout(k_rope, input_layout).contiguous().to(device)
49+ else:
50+ query_rope = None
51+ key_rope = None
52+ if pse is not None:
53+ pse1 = pse.to(device)
54+ else:
55+ pse1 = None
56+ 
57+ if atten_mask is not None:
58+ atten_mask1 = atten_mask.to(device)
59+ else:
60+ atten_mask1 = None
61+ 
62+ torch.npu.manual_seed(seed)
63+ 
64+ out_all = torch_npu.npu_fusion_attention_v2(
65+ q1, k1, v1, n1, input_layout,
66+ # optional parameters
67+ pse=pse1,
68+ atten_mask=atten_mask1,
69+ actual_seq_qlen=actual_seq_qlen,
70+ actual_seq_kvlen=actual_seq_kvlen,
71+ query_rope=query_rope,
72+ key_rope=key_rope,
73+ scale=scale,
74+ keep_prob=keep_prob,
75+ pre_tokens=pre_tokens,
76+ next_tokens=next_tokens,
77+ prefix=prefix,
78+ pse_type=pse_type,
79+ q_start_idx=tuple([q_start_idx]) if q_start_idx is not None else None,
80+ kv_start_idx=tuple([kv_start_idx]) if kv_start_idx is not None else None,
81+ sparse_mode=sparse_mode)
82+ torch.npu.synchronize()
83+ 
84+ npu_out = out_all[0].cpu()
85+ npu_max = out_all[1].cpu()
86+ npu_sum = out_all[2].cpu()
87+ 
88+ return npu_out, npu_max, npu_sum
Aattention/flash_attention_score/tests/pytest/test_case.py+199-0
@@ -0,0 +1,199 @@
1+#!/usr/bin/python
2+# -*- coding: utf-8 -*-
3+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+# See LICENSE in the root of the software repository for the full text of the License.
10+# ======================================================================================================================
11+ 
12+import torch
13+ 
14+####### 参数说明 ########
15+# B:必选; batch_size,TND格式下可选
16+# N1:必选; head_num
17+# N2:可选; kv's head_num,支持GQA/MHA/MQA
18+# Sq:必选; query's sequence length;TND格式下可选
19+# Skv:可选; key&value's sequence length
20+# D:必选; 表示query&key&value的head_dim
21+# DV:可选; value's head_dim;设置该参数,value的head_dim以DV为准
22+# DRope:可选; query_rope&key_rope's head_dim,支持D=128,DRope=64
23+# input_layout:必选; 输入tensor的格式, [BNSD, BSH, BSND, SBH, TND]
24+# dtype:必选; 数据类型, [torch.float16, torch.bfloat16, torch.float]
25+# scale:可选; 注意力得分缩放系数
26+# actual_seq_qlen:可选; TND下必选;query实际的序列长度
27+# actual_seq_kvlen:可选; TND下必选;key&value实际的序列长度
28+ 
29+# keep_prob:可选; dropout的保留概率：keep_prob = 1 - dropout_p;dropout的cpu实现方式基于Ascend 950PR/Ascend 950DT
30+# seed:可选; 随机种子，用于随机数生成
31+ 
32+# sparse_mode:可选; sparse模式, [0, 1, 2, 3, 4, 5, 6, 7, 8]
33+# - sparse_mode=3:等价于gpu的"causal=true"
34+# pre_tokens:可选; 配合sparse_mode使用,详细资料可参考算子文档
35+# next_tokens:可选; 配合sparse_mode使用,详细资料可参考算子文档
36+# prefix:可选; list类型,sparse_mode=5/6时需要设置,满足len(prefix)==batch_size
37+ 
38+# pse_type:可选; 位置编码, [0, 1, 2, 3]
39+# - 0/1: pse_layout取值在[bnss, 1nss, bn1s, bnhs, 1nhs],TND格式下取值在[bnhs, 1nhs]
40+# - 2: pse_layout取值在[bn, n],等价于gpu的"alibi=true"
41+# - 3: pse_layout取值在[bn, n]
42+# pse_layout:可选; 用于生成golden,取值范围[bnss, 1nss, bn1s, bnhs, 1nhs, bn, n],约束如上
43+# q_start_idx:可选; query的位置起始偏移
44+# kv_start_idx:可选; key&value的位置起始偏移
45+ 
46+TestCases = {
47+ "GQA_01": {
48+ "B": 4,
49+ "N1": 16,
50+ "N2": 8,
51+ "Sq": 256,
52+ "Skv": 256,
53+ "D": 128,
54+ "input_layout": "BSND",
55+ "dtype": torch.bfloat16,
56+ "sparse_mode": 3, # causal=true
57+ },
58+ "MLA_02": {
59+ "N1": 16,
60+ "N2": 8,
61+ "actual_seq_qlen": [256],
62+ "D": 128,
63+ "DRope": 64, # rope head-dim
64+ "input_layout": "TND",
65+ "dtype": torch.bfloat16,
66+ "sparse_mode": 3,
67+ },
68+ "MLA_03": {
69+ "B": 4,
70+ "N1": 16,
71+ "N2": 8,
72+ "Sq": 256,
73+ "D": 192, # query&key's head-dim
74+ "DV": 128, # value's head-dim
75+ "input_layout": "BSH",
76+ "dtype": torch.float16,
77+ "sparse_mode": 3,
78+ },
79+ "ALIBI_04": {
80+ "B": 4,
81+ "N1": 8,
82+ "Sq": 256,
83+ "D": 256, # query&key's head-dim
84+ "input_layout": "BNSD",
85+ "dtype": torch.bfloat16,
86+ "sparse_mode": 3,
87+ "pse_type": 2,
88+ "pse_layout": "bn" # for generate golden data
89+ },
90+ "ALIBI_05": {
91+ "B": 4,
92+ "N1": 8,
93+ "Sq": 256,
94+ "D": 128, # query&key's head-dim
95+ "input_layout": "BSND",
96+ "dtype": torch.float16,
97+ "sparse_mode": 3,
98+ "pse_type": 3, # special alibi
99+ "pse_layout": "n" # for generate golden data
100+ },
101+ "SPARSE_06": {
102+ "B": 4,
103+ "N1": 8,
104+ "Sq": 256,
105+ "D": 768, # query&key's head-dim
106+ "input_layout": "SBH",
107+ "dtype": torch.bfloat16,
108+ "sparse_mode": 0,
109+ "pre_tokens": 128,
110+ "next_tokens": 128,
111+ "pse_type": 3, # special alibi
112+ "pse_layout": "n" # for generate golden data
113+ },
114+ "SPARSE_07": {
115+ "B": 4,
116+ "N1": 8,
117+ "Sq": 256,
118+ "D": 128, # query&key's head-dim
119+ "input_layout": "BSND",
120+ "dtype": torch.float32,
121+ "sparse_mode": 4, # if sparse_mode=4, window_size=(pre_tokens, next_tokens)
122+ "pre_tokens": 128,
123+ "next_tokens": 128,
124+ "pse_type": 3, # special alibi
125+ "pse_layout": "n" # for generate golden data
126+ },
127+ "SPARSE_08": {
128+ "B": 4,
129+ "N1": 8,
130+ "Sq": 256,
131+ "D": 128, # query&key's head-dim
132+ "input_layout": "BSND",
133+ "dtype": torch.bfloat16,
134+ "sparse_mode": 5,
135+ "prefix": [100, 128, 130, 150],
136+ "pse_type": 3, # special alibi
137+ "pse_layout": "n" # for generate golden data
138+ },
139+ "DROPOUT_09": {
140+ "B": 4,
141+ "N1": 8,
142+ "Sq": 256,
143+ "Skv": 512,
144+ "D": 128, # query&key's head-dim
145+ "input_layout": "BSND",
146+ "dtype": torch.float32,
147+ "sparse_mode": 5,
148+ "prefix": [100, 128, 401, 300],
149+ "pse_type": 0, # special alibi
150+ "pse_layout": "bnss", # for generate golden data
151+ "keep_prob": 0.9, # keep_prob = 1 - dropout_p
152+ },
153+ "TND_01": {
154+ "N1": 8,
155+ "N2": 4,
156+ "actual_seq_qlen": [128, 256, 512],
157+ "D": 128, # query&key's head-dim
158+ "DRope": 64, # rope head-dim
159+ "input_layout": "TND",
160+ "dtype": torch.bfloat16,
161+ "sparse_mode": 3,
162+ },
163+ "TND_02": {
164+ "N1": 8,
165+ "actual_seq_qlen": [2048],
166+ "D": 192, # query&key's head-dim
167+ "input_layout": "TND",
168+ "dtype": torch.bfloat16,
169+ "sparse_mode": 6,
170+ "prefix": [1568]
171+ },
172+ "TND_03": {
173+ "N1": 8,
174+ "N2": 4,
175+ "actual_seq_qlen": [256, 512, 768],
176+ "actual_seq_kvlen": [256, 512, 768],
177+ "D": 128, # query&key's head-dim
178+ "input_layout": "TND",
179+ "dtype": torch.bfloat16,
180+ "sparse_mode": 3,
181+ "pse_type": 3, # special alibi
182+ "pse_layout": "n", # for generate golden data
183+ "keep_prob": 0.9, # keep_prob = 1 - dropout_p
184+ },
185+ "TND_04": {
186+ "N1": 8,
187+ "N2": 4,
188+ "actual_seq_qlen": [128, 256, 512],
189+ "actual_seq_kvlen": [256, 512, 768],
190+ "D": 128, # query&key's head-dim
191+ "input_layout": "TND",
192+ "dtype": torch.bfloat16,
193+ "sparse_mode": 8,
194+ "next_tokens": 64,
195+ "q_start_idx": 64,
196+ "kv_start_idx": 32,
197+ "keep_prob": 0.9, # keep_prob = 1 - dropout_p
198+ }
199+}
Mattention/flash_attention_score_grad/op_kernel/arch35/flash_attention_score_grad_entry_regbase.h+45-29
Mattention/flash_attention_score_grad/op_kernel/arch35/flash_attention_score_grad_s1s2_bn2gs1s2_post_regbase.h+1-1
Aattention/flash_attention_score_grad/tests/pytest/cpu_impl.py+1984-0文件内容审核中，请稍后刷新重试
Aattention/flash_attention_score_grad/tests/pytest/test_case.py+179-0文件内容审核中，请稍后刷新重试