Star1061
2569
代码
代码
Issues186
Pull Requests439
Actions
项目讨论
Wiki
分析
Star1061
2569
已合并
A5 BlockSparseAttention FP8场景性能改进 #6620
tangkaidi创建于 6月9日
A5 BlockSparseAttention FP8场景性能改进 #6620
已合并
tangkaidi创建于 6月9日
从已删除 :master合入到cann/ops-transformermaster
共 5 个文件变更+88-240
Mattention/block_sparse_attention/op_kernel/arch35/block_sparse_attention_kernel_arch35_full_quant.h+1-2
@@ -439,8 +439,7 @@ public:
439 curTileMod, gatheredKvSTileIdxDe,439 curTileMod, gatheredKvSTileIdxDe,
440 (gatheredKvSTileIdxDe == 0),440 (gatheredKvSTileIdxDe == 0),
441 (gatheredKvSTileIdxDe == kvSLoopNum - 1),441 (gatheredKvSTileIdxDe == kvSLoopNum - 1),
442- mm2ToReFlag,442+ mm2ToReFlag);
443- true);
444#endif443#endif
445 }444 }
446 }445 }
Mattention/block_sparse_attention/op_kernel/attn_infra/epilogue/block/block_epilogue_online_softmax_arch35_reg_low_prec.hpp+29-85
@@ -526,13 +526,6 @@ private:
526 AscendC::RoundMode::CAST_RINT,526 AscendC::RoundMode::CAST_RINT,
527 };527 };
528 528 
529- constexpr static CastTrait castTraitOneRINT = {
530- RegLayout::ONE,
531- SatMode::SAT,
532- MaskMergeMode::ZEROING,
533- AscendC::RoundMode::CAST_RINT,
534- };
535- 
536 constexpr static CastTrait castTraitTwoRINT = {529 constexpr static CastTrait castTraitTwoRINT = {
537 RegLayout::TWO,530 RegLayout::TWO,
538 SatMode::SAT,531 SatMode::SAT,
@@ -540,23 +533,12 @@ private:
540 AscendC::RoundMode::CAST_RINT,533 AscendC::RoundMode::CAST_RINT,
541 };534 };
542 535 
543- constexpr static CastTrait castTraitThreeRINT = {
544- RegLayout::THREE,
545- SatMode::SAT,
546- MaskMergeMode::ZEROING,
547- AscendC::RoundMode::CAST_RINT,
548- };
549- 
550 RegTensor<float> floatExpVreg0;536 RegTensor<float> floatExpVreg0;
551 RegTensor<float> floatExpVreg1;537 RegTensor<float> floatExpVreg1;
552- RegTensor<float> deInterleaveVreg0;
553- RegTensor<float> deInterleaveVreg1;
554- RegTensor<float> deInterleaveVreg2;
555- RegTensor<float> deInterleaveVreg3;
556 RegTensor<ElementOutput> pVreg0;538 RegTensor<ElementOutput> pVreg0;
557 RegTensor<ElementOutput> pVreg1;539 RegTensor<ElementOutput> pVreg1;
558- RegTensor<ElementOutput> pVreg2;540+ RegTensor<ElementOutput> deInterleaveVreg0;
559- RegTensor<ElementOutput> pVreg3;541+ RegTensor<ElementOutput> deInterleaveVreg1;
560 542 
561 MaskReg pRegUint8All = CreateMask<uint8_t, MaskPattern::ALL>();543 MaskReg pRegUint8All = CreateMask<uint8_t, MaskPattern::ALL>();
562 MaskReg pRegUint8VL128 = CreateMask<uint8_t, MaskPattern::VL128>();544 MaskReg pRegUint8VL128 = CreateMask<uint8_t, MaskPattern::VL128>();
@@ -577,19 +559,13 @@ private:
577 Cast<float, ElementInput, castTraitZeroUNKNOWN>(floatExpVreg0, expDstVreg, pRegFp16All);559 Cast<float, ElementInput, castTraitZeroUNKNOWN>(floatExpVreg0, expDstVreg, pRegFp16All);
578 Cast<float, ElementInput, castTraitOneUNKNOWN>(floatExpVreg1, expDstVreg, pRegFp16All);560 Cast<float, ElementInput, castTraitOneUNKNOWN>(floatExpVreg1, expDstVreg, pRegFp16All);
579 561 
580- DeInterleave(deInterleaveVreg0, deInterleaveVreg1, floatExpVreg0, floatExpVreg0);562+ Cast<ElementOutput, float, castTraitZeroRINT>(pVreg0, floatExpVreg0, pRegFp32All);
581- DeInterleave(deInterleaveVreg2, deInterleaveVreg3, floatExpVreg1, floatExpVreg1);563+ Cast<ElementOutput, float, castTraitTwoRINT>(pVreg1, floatExpVreg1, pRegFp32All);
582- 
583- Cast<ElementOutput, float, castTraitZeroRINT>(pVreg0, deInterleaveVreg0, pRegFp32All);
584- Cast<ElementOutput, float, castTraitOneRINT>(pVreg1, deInterleaveVreg2, pRegFp32All);
585- Cast<ElementOutput, float, castTraitTwoRINT>(pVreg2, deInterleaveVreg1, pRegFp32All);
586- Cast<ElementOutput, float, castTraitThreeRINT>(pVreg3, deInterleaveVreg3, pRegFp32All);
587 564 
588 Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg1, pRegUint8All);565 Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg1, pRegUint8All);
589- Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg2, pRegUint8All);566+ DeInterleave(deInterleaveVreg0, deInterleaveVreg1, pVreg0, pVreg0);
590- Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg3, pRegUint8All);
591 567 
592- StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(expUb + i * ELE_NUM_PER_C0_FP8, pVreg0,568+ StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(expUb + i * ELE_NUM_PER_C0_FP8, deInterleaveVreg0,
593 blockStride, pRegUint8VL128);569 blockStride, pRegUint8VL128);
594 ReduceSum(expSumVreg, expDstVregResult, pregFp16TailN);570 ReduceSum(expSumVreg, expDstVregResult, pregFp16TailN);
595 StoreUnAlign<ElementInput, PostLiteral::POST_MODE_UPDATE>(expSumUb, expSumVreg, expSumUreg, 1);571 StoreUnAlign<ElementInput, PostLiteral::POST_MODE_UPDATE>(expSumUb, expSumVreg, expSumUreg, 1);
@@ -636,13 +612,6 @@ private:
636 AscendC::RoundMode::CAST_RINT,612 AscendC::RoundMode::CAST_RINT,
637 };613 };
638 614 
639- constexpr static CastTrait castTraitOneRINT = {
640- RegLayout::ONE,
641- SatMode::SAT,
642- MaskMergeMode::ZEROING,
643- AscendC::RoundMode::CAST_RINT,
644- };
645- 
646 constexpr static CastTrait castTraitTwoRINT = {615 constexpr static CastTrait castTraitTwoRINT = {
647 RegLayout::TWO,616 RegLayout::TWO,
648 SatMode::SAT,617 SatMode::SAT,
@@ -650,35 +619,20 @@ private:
650 AscendC::RoundMode::CAST_RINT,619 AscendC::RoundMode::CAST_RINT,
651 };620 };
652 621 
653- constexpr static CastTrait castTraitThreeRINT = {
654- RegLayout::THREE,
655- SatMode::SAT,
656- MaskMergeMode::ZEROING,
657- AscendC::RoundMode::CAST_RINT,
658- };
659- 
660 RegTensor<ElementInput> expDstVreg0Fp8;622 RegTensor<ElementInput> expDstVreg0Fp8;
661 RegTensor<ElementInput> expDstVreg1Fp8;623 RegTensor<ElementInput> expDstVreg1Fp8;
662- RegTensor<float> floatVreg0;624+ RegTensor<float> floatExpVreg0;
663- RegTensor<float> floatVreg1;625+ RegTensor<float> floatExpVreg1;
664- RegTensor<float> floatVreg2;626+ RegTensor<float> floatExpVreg2;
665- RegTensor<float> floatVreg3;627+ RegTensor<float> floatExpVreg3;
666- RegTensor<float> deInterleaveVreg0;
667- RegTensor<float> deInterleaveVreg1;
668- RegTensor<float> deInterleaveVreg2;
669- RegTensor<float> deInterleaveVreg3;
670- RegTensor<float> deInterleaveVreg4;
671- RegTensor<float> deInterleaveVreg5;
672- RegTensor<float> deInterleaveVreg6;
673- RegTensor<float> deInterleaveVreg7;
674 RegTensor<ElementOutput> pVreg0;628 RegTensor<ElementOutput> pVreg0;
675 RegTensor<ElementOutput> pVreg1;629 RegTensor<ElementOutput> pVreg1;
676 RegTensor<ElementOutput> pVreg2;630 RegTensor<ElementOutput> pVreg2;
677 RegTensor<ElementOutput> pVreg3;631 RegTensor<ElementOutput> pVreg3;
678- RegTensor<ElementOutput> pVreg4;632+ RegTensor<ElementOutput> deInterleaveVreg0;
679- RegTensor<ElementOutput> pVreg5;633+ RegTensor<ElementOutput> deInterleaveVreg1;
680- RegTensor<ElementOutput> pVreg6;634+ RegTensor<ElementOutput> deInterleaveVreg2;
681- RegTensor<ElementOutput> pVreg7;635+ RegTensor<ElementOutput> deInterleaveVreg3;
682 636 
683 MaskReg pRegUint8All = CreateMask<uint8_t, MaskPattern::ALL>();637 MaskReg pRegUint8All = CreateMask<uint8_t, MaskPattern::ALL>();
684 MaskReg pRegUint8VL128 = CreateMask<uint8_t, MaskPattern::VL128>();638 MaskReg pRegUint8VL128 = CreateMask<uint8_t, MaskPattern::VL128>();
@@ -701,36 +655,26 @@ private:
701 Muls(expDstVreg0Fp8, expDstVreg0Fp8, maxValueFP8, pRegFp16All);655 Muls(expDstVreg0Fp8, expDstVreg0Fp8, maxValueFP8, pRegFp16All);
702 Muls(expDstVreg1Fp8, expDstVreg1Fp8, maxValueFP8, pregFp16TailN);656 Muls(expDstVreg1Fp8, expDstVreg1Fp8, maxValueFP8, pregFp16TailN);
703 657 
704- Cast<float, ElementInput, castTraitZeroUNKNOWN>(floatVreg0, expDstVreg0Fp8, pRegFp16All);658+ Cast<float, ElementInput, castTraitZeroUNKNOWN>(floatExpVreg0, expDstVreg0Fp8, pRegFp16All);
liwenguihw
liwenguihwliwenguihw6月9日

image.pngvf是否要考虑下划线代替驼峰

likedislike
T
tangkaidi
6月9日 评论：
705- Cast<float, ElementInput, castTraitOneUNKNOWN>(floatVreg1, expDstVreg0Fp8, pRegFp16All);659+ Cast<float, ElementInput, castTraitOneUNKNOWN>(floatExpVreg1, expDstVreg0Fp8, pRegFp16All);
706- Cast<float, ElementInput, castTraitZeroUNKNOWN>(floatVreg2, expDstVreg1Fp8, pRegFp16All);660+ Cast<float, ElementInput, castTraitZeroUNKNOWN>(floatExpVreg2, expDstVreg1Fp8, pRegFp16All);
707- Cast<float, ElementInput, castTraitOneUNKNOWN>(floatVreg3, expDstVreg1Fp8, pRegFp16All);661+ Cast<float, ElementInput, castTraitOneUNKNOWN>(floatExpVreg3, expDstVreg1Fp8, pRegFp16All);
708 662 
709- DeInterleave(deInterleaveVreg0, deInterleaveVreg1, floatVreg0, floatVreg0);663+ Cast<ElementOutput, float, castTraitZeroRINT>(pVreg0, floatExpVreg0, pRegFp32All);
710- DeInterleave(deInterleaveVreg2, deInterleaveVreg3, floatVreg1, floatVreg1);664+ Cast<ElementOutput, float, castTraitTwoRINT>(pVreg1, floatExpVreg1, pRegFp32All);
711- DeInterleave(deInterleaveVreg4, deInterleaveVreg5, floatVreg2, floatVreg2);665+ Cast<ElementOutput, float, castTraitZeroRINT>(pVreg2, floatExpVreg2, pRegFp32All);
712- DeInterleave(deInterleaveVreg6, deInterleaveVreg7, floatVreg3, floatVreg3);666+ Cast<ElementOutput, float, castTraitTwoRINT>(pVreg3, floatExpVreg3, pRegFp32All);
713- 
714- Cast<ElementOutput, float, castTraitZeroRINT>(pVreg0, deInterleaveVreg0, pRegFp32All);
715- Cast<ElementOutput, float, castTraitOneRINT>(pVreg1, deInterleaveVreg2, pRegFp32All);
716- Cast<ElementOutput, float, castTraitTwoRINT>(pVreg2, deInterleaveVreg1, pRegFp32All);
717- Cast<ElementOutput, float, castTraitThreeRINT>(pVreg3, deInterleaveVreg3, pRegFp32All);
718- Cast<ElementOutput, float, castTraitZeroRINT>(pVreg4, deInterleaveVreg4, pRegFp32All);
719- Cast<ElementOutput, float, castTraitOneRINT>(pVreg5, deInterleaveVreg6, pRegFp32All);
720- Cast<ElementOutput, float, castTraitTwoRINT>(pVreg6, deInterleaveVreg5, pRegFp32All);
721- Cast<ElementOutput, float, castTraitThreeRINT>(pVreg7, deInterleaveVreg7, pRegFp32All);
722 667 
723 Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg1, pRegUint8All);668 Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg1, pRegUint8All);
724- Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg2, pRegUint8All);669+ Or((RegTensor<uint8_t> &)pVreg2, (RegTensor<uint8_t> &)pVreg2, (RegTensor<uint8_t> &)pVreg3, pRegUint8All);
725- Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg3, pRegUint8All);670+ DeInterleave(deInterleaveVreg0, deInterleaveVreg1, pVreg0, pVreg0);
726- Or((RegTensor<uint8_t> &)pVreg4, (RegTensor<uint8_t> &)pVreg4, (RegTensor<uint8_t> &)pVreg5, pRegUint8All);671+ DeInterleave(deInterleaveVreg2, deInterleaveVreg3, pVreg2, pVreg2);
727- Or((RegTensor<uint8_t> &)pVreg4, (RegTensor<uint8_t> &)pVreg4, (RegTensor<uint8_t> &)pVreg6, pRegUint8All);
728- Or((RegTensor<uint8_t> &)pVreg4, (RegTensor<uint8_t> &)pVreg4, (RegTensor<uint8_t> &)pVreg7, pRegUint8All);
729 672 
730- StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(expUb + i * ELE_NUM_PER_C0_FP8, pVreg0,673+ StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(expUb + i * ELE_NUM_PER_C0_FP8, deInterleaveVreg0,
731 blockStride, pRegUint8VL128);674 blockStride, pRegUint8VL128);
732- StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(675+ StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(expUb + i * ELE_NUM_PER_C0_FP8 +
733- expUb + i * ELE_NUM_PER_C0_FP8 + blockStride * HALF_VECTOR_SIZE, pVreg4, blockStride, pRegUint8VL128);676+ blockStride * HALF_VECTOR_SIZE,
677+ deInterleaveVreg2, blockStride, pRegUint8VL128);
734 Add<ElementInput, MaskMergeMode::MERGING>(expDstVreg0, expDstVreg1, expDstVreg0, pregFp16TailN);678 Add<ElementInput, MaskMergeMode::MERGING>(expDstVreg0, expDstVreg1, expDstVreg0, pregFp16TailN);
735 ReduceSum(expSumVreg, expDstVreg0, pRegFp16All);679 ReduceSum(expSumVreg, expDstVreg0, pRegFp16All);
736 StoreUnAlign<ElementInput, PostLiteral::POST_MODE_UPDATE>(expSumUb, expSumVreg, expSumUreg, 1);680 StoreUnAlign<ElementInput, PostLiteral::POST_MODE_UPDATE>(expSumUb, expSumVreg, expSumUreg, 1);
Mattention/block_sparse_attention/op_kernel/attn_infra/epilogue/block/block_epilogue_online_softmax_arch35_reg_low_prec_bf16.hpp+33-88
@@ -696,13 +696,6 @@ private:
696 AscendC::RoundMode::CAST_RINT,696 AscendC::RoundMode::CAST_RINT,
697 };697 };
698 698 
699- constexpr static CastTrait castTraitOneRINT = {
700- RegLayout::ONE,
701- SatMode::SAT,
702- MaskMergeMode::ZEROING,
703- AscendC::RoundMode::CAST_RINT,
704- };
705- 
706 constexpr static CastTrait castTraitTwoRINT = {699 constexpr static CastTrait castTraitTwoRINT = {
707 RegLayout::TWO,700 RegLayout::TWO,
708 SatMode::SAT,701 SatMode::SAT,
@@ -710,25 +703,14 @@ private:
710 AscendC::RoundMode::CAST_RINT,703 AscendC::RoundMode::CAST_RINT,
711 };704 };
712 705 
713- constexpr static CastTrait castTraitThreeRINT = {
714- RegLayout::THREE,
715- SatMode::SAT,
716- MaskMergeMode::ZEROING,
717- AscendC::RoundMode::CAST_RINT,
718- };
719- 
720- RegTensor<float> deInterleaveVreg0;
721- RegTensor<float> deInterleaveVreg1;
722- RegTensor<float> deInterleaveVreg2;
723- RegTensor<float> deInterleaveVreg3;
724 RegTensor<ElementOutput> pVreg0;706 RegTensor<ElementOutput> pVreg0;
725 RegTensor<ElementOutput> pVreg1;707 RegTensor<ElementOutput> pVreg1;
726- RegTensor<ElementOutput> pVreg2;708+ RegTensor<ElementOutput> deInterleaveVreg0;
727- RegTensor<ElementOutput> pVreg3;709+ RegTensor<ElementOutput> deInterleaveVreg1;
728 710 
729 MaskReg pRegUint8All = CreateMask<uint8_t, MaskPattern::ALL>();711 MaskReg pRegUint8All = CreateMask<uint8_t, MaskPattern::ALL>();
730 MaskReg pRegUint8VL128 = CreateMask<uint8_t, MaskPattern::VL128>();712 MaskReg pRegUint8VL128 = CreateMask<uint8_t, MaskPattern::VL128>();
731- MaskReg pRegFp16All = CreateMask<ElementInput, MaskPattern::ALL>();713+ MaskReg pRegBf16All = CreateMask<ElementInput, MaskPattern::ALL>();
732 MaskReg pRegFp32All = CreateMask<float, MaskPattern::ALL>();714 MaskReg pRegFp32All = CreateMask<float, MaskPattern::ALL>();
733 MaskReg pregFp32tailNOdd = UpdateMask<float>(tailNOdd);715 MaskReg pregFp32tailNOdd = UpdateMask<float>(tailNOdd);
734 MaskReg pregFp32tailNEven = UpdateMask<float>(tailNEven);716 MaskReg pregFp32tailNEven = UpdateMask<float>(tailNEven);
@@ -737,8 +719,8 @@ private:
737 LoadAlign<float, LoadDist::DIST_BRC_B32>(maxVreg, nowMaxUb + i);719 LoadAlign<float, LoadDist::DIST_BRC_B32>(maxVreg, nowMaxUb + i);
738 Duplicate(expSumVreg, 0);720 Duplicate(expSumVreg, 0);
739 LoadAlign(expVreg, srcUb + i * S2BaseSize);721 LoadAlign(expVreg, srcUb + i * S2BaseSize);
740- Cast<float, ElementInput, castTraitZeroUNKNOWN>(expFloatVreg0, expVreg, pRegFp16All);722+ Cast<float, ElementInput, castTraitZeroUNKNOWN>(expFloatVreg0, expVreg, pRegBf16All);
741- Cast<float, ElementInput, castTraitOneUNKNOWN>(expFloatVreg1, expVreg, pRegFp16All);723+ Cast<float, ElementInput, castTraitOneUNKNOWN>(expFloatVreg1, expVreg, pRegBf16All);
742 FusedExpSub(expDstFloatVreg0, expFloatVreg0, maxVreg, pregFp32tailNEven);724 FusedExpSub(expDstFloatVreg0, expFloatVreg0, maxVreg, pregFp32tailNEven);
743 FusedExpSub(expDstFloatVreg1, expFloatVreg1, maxVreg, pregFp32tailNOdd);725 FusedExpSub(expDstFloatVreg1, expFloatVreg1, maxVreg, pregFp32tailNOdd);
744 Add<float, MaskMergeMode::MERGING>(expSumVreg, expSumVreg, expDstFloatVreg0, pregFp32tailNEven);726 Add<float, MaskMergeMode::MERGING>(expSumVreg, expSumVreg, expDstFloatVreg0, pregFp32tailNEven);
@@ -748,21 +730,15 @@ private:
748 Muls(expDstFloatVreg0, expDstFloatVreg0, maxValueFP8, pregFp32tailNEven);730 Muls(expDstFloatVreg0, expDstFloatVreg0, maxValueFP8, pregFp32tailNEven);
749 Muls(expDstFloatVreg1, expDstFloatVreg1, maxValueFP8, pregFp32tailNOdd);731 Muls(expDstFloatVreg1, expDstFloatVreg1, maxValueFP8, pregFp32tailNOdd);
750 732 
751- DeInterleave(deInterleaveVreg0, deInterleaveVreg1, expDstFloatVreg0, expDstFloatVreg0);733+ Cast<ElementOutput, float, castTraitZeroRINT>(pVreg0, expDstFloatVreg0, pRegFp32All);
752- DeInterleave(deInterleaveVreg2, deInterleaveVreg3, expDstFloatVreg1, expDstFloatVreg1);734+ Cast<ElementOutput, float, castTraitTwoRINT>(pVreg1, expDstFloatVreg1, pRegFp32All);
753- 
754- Cast<ElementOutput, float, castTraitZeroRINT>(pVreg0, deInterleaveVreg0, pRegFp32All);
755- Cast<ElementOutput, float, castTraitOneRINT>(pVreg1, deInterleaveVreg2, pRegFp32All);
756- Cast<ElementOutput, float, castTraitTwoRINT>(pVreg2, deInterleaveVreg1, pRegFp32All);
757- Cast<ElementOutput, float, castTraitThreeRINT>(pVreg3, deInterleaveVreg3, pRegFp32All);
758 735 
759 Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg1, pRegUint8All);736 Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg1, pRegUint8All);
760- Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg2, pRegUint8All);737+ DeInterleave(deInterleaveVreg0, deInterleaveVreg1, pVreg0, pVreg0);
761- Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg3, pRegUint8All);
762 738 
763- StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(expUb + i * ELE_NUM_PER_C0_FP8, pVreg0,739+ StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(expUb + i * ELE_NUM_PER_C0_FP8, deInterleaveVreg0,
764 blockStride, pRegUint8VL128);740 blockStride, pRegUint8VL128);
765- ReduceSum(expSumVreg, expSumVreg, pRegFp16All);741+ ReduceSum(expSumVreg, expSumVreg, pRegBf16All);
766 StoreUnAlign<float, PostLiteral::POST_MODE_UPDATE>(expSumUb, expSumVreg, expSumUreg, 1);742 StoreUnAlign<float, PostLiteral::POST_MODE_UPDATE>(expSumUb, expSumVreg, expSumUreg, 1);
767 }743 }
768 vstas(expSumUreg, expSumUb, 0, POST_UPDATE);744 vstas(expSumUreg, expSumUb, 0, POST_UPDATE);
@@ -816,13 +792,6 @@ private:
816 AscendC::RoundMode::CAST_RINT,792 AscendC::RoundMode::CAST_RINT,
817 };793 };
818 794 
819- constexpr static CastTrait castTraitOneRINT = {
820- RegLayout::ONE,
821- SatMode::SAT,
822- MaskMergeMode::ZEROING,
823- AscendC::RoundMode::CAST_RINT,
824- };
825- 
826 constexpr static CastTrait castTraitTwoRINT = {795 constexpr static CastTrait castTraitTwoRINT = {
827 RegLayout::TWO,796 RegLayout::TWO,
828 SatMode::SAT,797 SatMode::SAT,
@@ -830,34 +799,19 @@ private:
830 AscendC::RoundMode::CAST_RINT,799 AscendC::RoundMode::CAST_RINT,
831 };800 };
832 801 
833- constexpr static CastTrait castTraitThreeRINT = {
834- RegLayout::THREE,
835- SatMode::SAT,
836- MaskMergeMode::ZEROING,
837- AscendC::RoundMode::CAST_RINT,
838- };
839- 
840- RegTensor<float> deInterleaveVreg0;
841- RegTensor<float> deInterleaveVreg1;
842- RegTensor<float> deInterleaveVreg2;
843- RegTensor<float> deInterleaveVreg3;
844- RegTensor<float> deInterleaveVreg4;
845- RegTensor<float> deInterleaveVreg5;
846- RegTensor<float> deInterleaveVreg6;
847- RegTensor<float> deInterleaveVreg7;
848 RegTensor<ElementOutput> pVreg0;802 RegTensor<ElementOutput> pVreg0;
849 RegTensor<ElementOutput> pVreg1;803 RegTensor<ElementOutput> pVreg1;
850 RegTensor<ElementOutput> pVreg2;804 RegTensor<ElementOutput> pVreg2;
851 RegTensor<ElementOutput> pVreg3;805 RegTensor<ElementOutput> pVreg3;
852- RegTensor<ElementOutput> pVreg4;806+ RegTensor<ElementOutput> deInterleaveVreg0;
853- RegTensor<ElementOutput> pVreg5;807+ RegTensor<ElementOutput> deInterleaveVreg1;
854- RegTensor<ElementOutput> pVreg6;808+ RegTensor<ElementOutput> deInterleaveVreg2;
855- RegTensor<ElementOutput> pVreg7;809+ RegTensor<ElementOutput> deInterleaveVreg3;
856 810 
857 MaskReg pRegUint8All = CreateMask<uint8_t, MaskPattern::ALL>();811 MaskReg pRegUint8All = CreateMask<uint8_t, MaskPattern::ALL>();
858 MaskReg pRegUint8VL128 = CreateMask<uint8_t, MaskPattern::VL128>();812 MaskReg pRegUint8VL128 = CreateMask<uint8_t, MaskPattern::VL128>();
859- MaskReg pRegFp16All = CreateMask<ElementInput, MaskPattern::ALL>();813+ MaskReg pRegBf16All = CreateMask<ElementInput, MaskPattern::ALL>();
860- MaskReg pregFp16TailN = UpdateMask<ElementInput>(tailN);814+ MaskReg pRegBf16TailN = UpdateMask<ElementInput>(tailN);
861 MaskReg pRegFp32All = CreateMask<float, MaskPattern::ALL>();815 MaskReg pRegFp32All = CreateMask<float, MaskPattern::ALL>();
862 MaskReg pregFp32TailNOdd = UpdateMask<float>(tailNOdd);816 MaskReg pregFp32TailNOdd = UpdateMask<float>(tailNOdd);
863 MaskReg pregFp32tailNEven = UpdateMask<float>(tailNEven);817 MaskReg pregFp32tailNEven = UpdateMask<float>(tailNEven);
@@ -867,10 +821,10 @@ private:
867 Duplicate(expSumVreg, 0);821 Duplicate(expSumVreg, 0);
868 LoadAlign(expVreg0, srcUb + i * S2BaseSize);822 LoadAlign(expVreg0, srcUb + i * S2BaseSize);
869 LoadAlign(expVreg1, srcUb + i * S2BaseSize + HALF_REP_SIZE);823 LoadAlign(expVreg1, srcUb + i * S2BaseSize + HALF_REP_SIZE);
870- Cast<float, ElementInput, castTraitZeroUNKNOWN>(expFloatVreg0, expVreg0, pRegFp16All);824+ Cast<float, ElementInput, castTraitZeroUNKNOWN>(expFloatVreg0, expVreg0, pRegBf16All);
871- Cast<float, ElementInput, castTraitOneUNKNOWN>(expFloatVreg1, expVreg0, pRegFp16All);825+ Cast<float, ElementInput, castTraitOneUNKNOWN>(expFloatVreg1, expVreg0, pRegBf16All);
872- Cast<float, ElementInput, castTraitZeroUNKNOWN>(expFloatVreg2, expVreg1, pRegFp16All);826+ Cast<float, ElementInput, castTraitZeroUNKNOWN>(expFloatVreg2, expVreg1, pRegBf16All);
873- Cast<float, ElementInput, castTraitOneUNKNOWN>(expFloatVreg3, expVreg1, pRegFp16All);827+ Cast<float, ElementInput, castTraitOneUNKNOWN>(expFloatVreg3, expVreg1, pRegBf16All);
874 FusedExpSub(expDstFloatVreg0, expFloatVreg0, maxVreg, pRegFp32All);828 FusedExpSub(expDstFloatVreg0, expFloatVreg0, maxVreg, pRegFp32All);
875 FusedExpSub(expDstFloatVreg1, expFloatVreg1, maxVreg, pRegFp32All);829 FusedExpSub(expDstFloatVreg1, expFloatVreg1, maxVreg, pRegFp32All);
876 FusedExpSub(expDstFloatVreg2, expFloatVreg2, maxVreg, pregFp32tailNEven);830 FusedExpSub(expDstFloatVreg2, expFloatVreg2, maxVreg, pregFp32tailNEven);
@@ -886,31 +840,22 @@ private:
886 Muls(expDstFloatVreg2, expDstFloatVreg2, maxValueFP8, pregFp32tailNEven);840 Muls(expDstFloatVreg2, expDstFloatVreg2, maxValueFP8, pregFp32tailNEven);
887 Muls(expDstFloatVreg3, expDstFloatVreg3, maxValueFP8, pregFp32TailNOdd);841 Muls(expDstFloatVreg3, expDstFloatVreg3, maxValueFP8, pregFp32TailNOdd);
888 842 
889- DeInterleave(deInterleaveVreg0, deInterleaveVreg1, expDstFloatVreg0, expDstFloatVreg0);843+ Cast<ElementOutput, float, castTraitZeroRINT>(pVreg0, expDstFloatVreg0, pRegFp32All);
890- DeInterleave(deInterleaveVreg2, deInterleaveVreg3, expDstFloatVreg1, expDstFloatVreg1);844+ Cast<ElementOutput, float, castTraitTwoRINT>(pVreg1, expDstFloatVreg1, pRegFp32All);
891- DeInterleave(deInterleaveVreg4, deInterleaveVreg5, expDstFloatVreg2, expDstFloatVreg2);845+ Cast<ElementOutput, float, castTraitZeroRINT>(pVreg2, expDstFloatVreg2, pRegFp32All);
892- DeInterleave(deInterleaveVreg6, deInterleaveVreg7, expDstFloatVreg3, expDstFloatVreg3);846+ Cast<ElementOutput, float, castTraitTwoRINT>(pVreg3, expDstFloatVreg3, pRegFp32All);
893- 
894- Cast<ElementOutput, float, castTraitZeroRINT>(pVreg0, deInterleaveVreg0, pRegFp32All);
895- Cast<ElementOutput, float, castTraitOneRINT>(pVreg1, deInterleaveVreg2, pRegFp32All);
896- Cast<ElementOutput, float, castTraitTwoRINT>(pVreg2, deInterleaveVreg1, pRegFp32All);
897- Cast<ElementOutput, float, castTraitThreeRINT>(pVreg3, deInterleaveVreg3, pRegFp32All);
898- Cast<ElementOutput, float, castTraitZeroRINT>(pVreg4, deInterleaveVreg4, pRegFp32All);
899- Cast<ElementOutput, float, castTraitOneRINT>(pVreg5, deInterleaveVreg6, pRegFp32All);
900- Cast<ElementOutput, float, castTraitTwoRINT>(pVreg6, deInterleaveVreg5, pRegFp32All);
901- Cast<ElementOutput, float, castTraitThreeRINT>(pVreg7, deInterleaveVreg7, pRegFp32All);
902 847 
903 Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg1, pRegUint8All);848 Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg1, pRegUint8All);
904- Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg2, pRegUint8All);849+ Or((RegTensor<uint8_t> &)pVreg2, (RegTensor<uint8_t> &)pVreg2, (RegTensor<uint8_t> &)pVreg3, pRegUint8All);
905- Or((RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg0, (RegTensor<uint8_t> &)pVreg3, pRegUint8All);850+ DeInterleave(deInterleaveVreg0, deInterleaveVreg1, pVreg0, pVreg0);
906- Or((RegTensor<uint8_t> &)pVreg4, (RegTensor<uint8_t> &)pVreg4, (RegTensor<uint8_t> &)pVreg5, pRegUint8All);851+ DeInterleave(deInterleaveVreg2, deInterleaveVreg3, pVreg2, pVreg2);
907- Or((RegTensor<uint8_t> &)pVreg4, (RegTensor<uint8_t> &)pVreg4, (RegTensor<uint8_t> &)pVreg6, pRegUint8All);852+ 
908- Or((RegTensor<uint8_t> &)pVreg4, (RegTensor<uint8_t> &)pVreg4, (RegTensor<uint8_t> &)pVreg7, pRegUint8All);853+ StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(expUb + i * ELE_NUM_PER_C0_FP8, deInterleaveVreg0,
909- StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(expUb + i * ELE_NUM_PER_C0_FP8, pVreg0,
910 blockStride, pRegUint8VL128);854 blockStride, pRegUint8VL128);
911- StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(855+ StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(expUb + i * ELE_NUM_PER_C0_FP8 +
912- expUb + i * ELE_NUM_PER_C0_FP8 + blockStride * HALF_VECTOR_SIZE, pVreg4, blockStride, pRegUint8VL128);856+ blockStride * HALF_VECTOR_SIZE,
913- ReduceSum(expSumVreg, expSumVreg, pRegFp16All);857+ deInterleaveVreg2, blockStride, pRegUint8VL128);
858+ ReduceSum(expSumVreg, expSumVreg, pRegBf16All);
914 StoreUnAlign<float, PostLiteral::POST_MODE_UPDATE>(expSumUb, expSumVreg, expSumUreg, 1);859 StoreUnAlign<float, PostLiteral::POST_MODE_UPDATE>(expSumUb, expSumVreg, expSumUreg, 1);
915 }860 }
916 vstas(expSumUreg, expSumUb, 0, POST_UPDATE);861 vstas(expSumUreg, expSumUb, 0, POST_UPDATE);
Mattention/block_sparse_attention/op_kernel/attn_infra/epilogue/block/block_epilogue_rescale_o_arch35_reg_high_prec.hpp+23-63
@@ -31,6 +31,7 @@ template <
31 class ElementO_,31 class ElementO_,
32 class ElementOTmp_,32 class ElementOTmp_,
33 class ElementS_,33 class ElementS_,
34+ class ElementKV_,
34 class TileCopy_,35 class TileCopy_,
35 class OTmpSrcPos_, // the src TPosition of pv res, viable configurations: GM/L0C36 class OTmpSrcPos_, // the src TPosition of pv res, viable configurations: GM/L0C
36 LseMode LSE_MODE_,37 LseMode LSE_MODE_,
@@ -40,6 +41,7 @@ class BlockEpilogue<
40 ElementO_,41 ElementO_,
41 ElementOTmp_,42 ElementOTmp_,
42 ElementS_,43 ElementS_,
44+ ElementKV_,
43 TileCopy_,45 TileCopy_,
44 OTmpSrcPos_>46 OTmpSrcPos_>
45{47{
@@ -50,6 +52,7 @@ public:
50 using ElementOTmp = ElementOTmp_;52 using ElementOTmp = ElementOTmp_;
51 using ElementLse = float;53 using ElementLse = float;
52 using SMDtype = ElementS_;54 using SMDtype = ElementS_;
55+ using ElementKV = ElementKV_;
53 using TileCopy = TileCopy_;56 using TileCopy = TileCopy_;
54 using OTmpSrcPos = OTmpSrcPos_;57 using OTmpSrcPos = OTmpSrcPos_;
55 58 
@@ -61,7 +64,8 @@ public:
61 static constexpr uint32_t RESCALE_ROW_MAX_ELEM_NUM = 64;64 static constexpr uint32_t RESCALE_ROW_MAX_ELEM_NUM = 64;
62 static constexpr uint32_t RESCALE_COL_MAX_ELEM_NUM = 128;65 static constexpr uint32_t RESCALE_COL_MAX_ELEM_NUM = 128;
63 static constexpr uint32_t RESCALE_VREG_SIZE = 256 / sizeof(ElementOTmp);66 static constexpr uint32_t RESCALE_VREG_SIZE = 256 / sizeof(ElementOTmp);
64- 67+ static constexpr float MAX_VALUE_RECIPROCAL_FP8 = 1.0f / 448.0f;
68+ static constexpr bool FULL_QUANT_FP8 = AscendC::IsSameType<ElementKV, fp8_e4m3fn_t>::value;
65 __aicore__ inline69 __aicore__ inline
66 BlockEpilogue(Arch::Resource<ArchTag> &resource)70 BlockEpilogue(Arch::Resource<ArchTag> &resource)
67 {71 {
@@ -142,8 +146,7 @@ public:
142 bool isFirstKvSTile,146 bool isFirstKvSTile,
143 bool isLastKvSTile,147 bool isLastKvSTile,
144 uint32_t colStrideCurSubCore,148 uint32_t colStrideCurSubCore,
145- Arch::CrossCoreFlag mm2ToReFlag,149+ Arch::CrossCoreFlag mm2ToReFlag)
146- bool isFullQuantFp8)
147 {150 {
148 uint32_t rowNumCurSubCore = tla::get<0>(gOTensorTlaTile.shape());151 uint32_t rowNumCurSubCore = tla::get<0>(gOTensorTlaTile.shape());
149 uint32_t colNumCurSubCore = tla::get<1>(gOTensorTlaTile.shape());152 uint32_t colNumCurSubCore = tla::get<1>(gOTensorTlaTile.shape());
@@ -218,16 +221,7 @@ public:
218 CopyUbToGmLse(gLseTensorTlaTile, ubLseTensorTla);221 CopyUbToGmLse(gLseTensorTlaTile, ubLseTensorTla);
219 }222 }
220 AscendC::PipeBarrier<PIPE_V>();223 AscendC::PipeBarrier<PIPE_V>();
221- if (isFullQuantFp8) {224+ 
222- if (dStages == 1) {
223- deQuantScaleO<DRegSplitStages::ONE>(goUb, rowNumCurSubCore, colStrideCurSubCore, colTail,
224- vlElemNum);
225- } else if (dStages == 2) {
226- deQuantScaleO<DRegSplitStages::TWO>(goUb, rowNumCurSubCore, colStrideCurSubCore, colTail,
227- vlElemNum);
228- }
229- AscendC::PipeBarrier<PIPE_V>();
230- }
231 if (std::is_same<ElementO, bfloat16_t>::value) {225 if (std::is_same<ElementO, bfloat16_t>::value) {
232 AscendC::Cast(226 AscendC::Cast(
233 goUbTensor16, goUbTensor32,227 goUbTensor16, goUbTensor32,
@@ -351,6 +345,9 @@ public:
351 Mul(mulVreg, goPreVreg, dmVreg, pregTail);345 Mul(mulVreg, goPreVreg, dmVreg, pregTail);
352 Add(goCurVreg, mulVreg, loVreg, pregTail);346 Add(goCurVreg, mulVreg, loVreg, pregTail);
353 Div(divVreg, goCurVreg, glVreg, pregTail);347 Div(divVreg, goCurVreg, glVreg, pregTail);
348+ if constexpr (FULL_QUANT_FP8) {
349+ Muls(divVreg, divVreg, MAX_VALUE_RECIPROCAL_FP8, pregTail);
350+ }
354 StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride, divVreg, pregTail);351 StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride, divVreg, pregTail);
355 }352 }
356 }353 }
@@ -390,6 +387,10 @@ public:
390 Add(goCurVreg1, mulVreg1, loVreg1, pregTail);387 Add(goCurVreg1, mulVreg1, loVreg1, pregTail);
391 Div(divVreg0, goCurVreg0, glVreg, pregFull);388 Div(divVreg0, goCurVreg0, glVreg, pregFull);
392 Div(divVreg1, goCurVreg1, glVreg, pregTail);389 Div(divVreg1, goCurVreg1, glVreg, pregTail);
390+ if constexpr (FULL_QUANT_FP8) {
391+ Muls(divVreg0, divVreg0, MAX_VALUE_RECIPROCAL_FP8, pregFull);
392+ Muls(divVreg1, divVreg1, MAX_VALUE_RECIPROCAL_FP8, pregTail);
393+ }
393 StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride, divVreg0, pregFull);394 StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride, divVreg0, pregFull);
394 StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride + vlElemNum, divVreg1, pregTail);395 StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride + vlElemNum, divVreg1, pregTail);
395 }396 }
@@ -418,6 +419,9 @@ public:
418 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(glVreg, glUb + i);419 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(glVreg, glUb + i);
419 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goCurVreg, loUb + i * colStride);420 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goCurVreg, loUb + i * colStride);
420 Div(divVreg, goCurVreg, glVreg, pregTail);421 Div(divVreg, goCurVreg, glVreg, pregTail);
422+ if constexpr (FULL_QUANT_FP8) {
423+ Muls(divVreg, divVreg, MAX_VALUE_RECIPROCAL_FP8, pregTail);
424+ }
421 StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride, divVreg, pregTail);425 StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride, divVreg, pregTail);
422 }426 }
423 }427 }
@@ -442,56 +446,15 @@ public:
442 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goCurVreg1, loUb + i * colStride + vlElemNum);446 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goCurVreg1, loUb + i * colStride + vlElemNum);
443 Div(divVreg0, goCurVreg0, glVreg, pregFull);447 Div(divVreg0, goCurVreg0, glVreg, pregFull);
444 Div(divVreg1, goCurVreg1, glVreg, pregTail);448 Div(divVreg1, goCurVreg1, glVreg, pregTail);
449+ if constexpr (FULL_QUANT_FP8) {
450+ Muls(divVreg0, divVreg0, MAX_VALUE_RECIPROCAL_FP8, pregFull);
451+ Muls(divVreg1, divVreg1, MAX_VALUE_RECIPROCAL_FP8, pregTail);
452+ }
445 StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride, divVreg0, pregFull);453 StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride, divVreg0, pregFull);
446 StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride + vlElemNum, divVreg1, pregTail);454 StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride + vlElemNum, divVreg1, pregTail);
447 }455 }
448 }456 }
449 457 
450- template <DRegSplitStages dRegSplitStages>
451- __simd_vf__ inline void deQuantScaleO(__ubuf__ ElementOTmp *goUb, uint32_t row, uint32_t colStride,
452- uint32_t colTail, uint32_t vlElemNum)
453- {
454- }
455- 
456- template <>
457- __simd_vf__ inline void deQuantScaleO<DRegSplitStages::ONE>(__ubuf__ ElementOTmp *goUb, uint32_t row,
458- uint32_t colStride, uint32_t colTail,
459- uint32_t vlElemNum)
460- {
461- using namespace AscendC::MicroAPI;
462- RegTensor<ElementOTmp> oVreg;
463- MaskReg pregTail = UpdateMask<float>(colTail);
464- 
465- for (uint16_t i = 0; i < row; ++i) {
466- LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(oVreg, goUb + i * colStride);
467- constexpr float maxValueReciprocal = 1.0f / 448.0f;
468- Muls(oVreg, oVreg, maxValueReciprocal, pregTail);
469- StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride, oVreg, pregTail);
470- }
471- }
472- 
473- template <>
474- __simd_vf__ inline void deQuantScaleO<DRegSplitStages::TWO>(__ubuf__ ElementOTmp *goUb, uint32_t row,
475- uint32_t colStride, uint32_t colTail,
476- uint32_t vlElemNum)
477- {
478- using namespace AscendC::MicroAPI;
479- RegTensor<ElementOTmp> oVreg0;
480- RegTensor<ElementOTmp> oVreg1;
481- MaskReg pregAll = CreateMask<float, MaskPattern::ALL>();
482- MaskReg pregTail = UpdateMask<float>(colTail);
483- 
484- for (uint16_t i = 0; i < row; ++i) {
485- LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(oVreg0, goUb + i * colStride);
486- LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(oVreg1, goUb + i * colStride + vlElemNum);
487- constexpr float maxValueReciprocal = 1.0f / 448.0f;
488- Muls(oVreg0, oVreg0, maxValueReciprocal, pregAll);
489- Muls(oVreg1, oVreg1, maxValueReciprocal, pregTail);
490- StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride, oVreg0, pregAll);
491- StoreAlign<ElementOTmp, StoreDist::DIST_NORM_B32>(goUb + i * colStride + vlElemNum, oVreg1, pregTail);
492- }
493- }
494- 
495 // When qS can be copied from UB to GM continuously,458 // When qS can be copied from UB to GM continuously,
496 // qS would be stored from reg to UB continuously.459 // qS would be stored from reg to UB continuously.
497 __simd_vf__ inline void LogSumExpFuncQSAxisContinuous(460 __simd_vf__ inline void LogSumExpFuncQSAxisContinuous(
@@ -561,8 +524,7 @@ public:
561 uint32_t gatheredKvSTileIdx,524 uint32_t gatheredKvSTileIdx,
562 bool isFirstKvSTile,525 bool isFirstKvSTile,
563 bool isLastKvSTile,526 bool isLastKvSTile,
564- Arch::CrossCoreFlag mm2ToReFlag,527+ Arch::CrossCoreFlag mm2ToReFlag)
565- bool isFullQuantFp8 = false)
566 {528 {
567 uint32_t rowNumOri = actualOriShape[0];529 uint32_t rowNumOri = actualOriShape[0];
568 uint32_t colNumOri = actualOriShape[1];530 uint32_t colNumOri = actualOriShape[1];
@@ -594,8 +556,7 @@ public:
594 isFirstKvSTile,556 isFirstKvSTile,
595 isLastKvSTile,557 isLastKvSTile,
596 colStrideCurSubCore,558 colStrideCurSubCore,
597- mm2ToReFlag,559+ mm2ToReFlag);
598- isFullQuantFp8);
599 } else {560 } else {
600 Arch::CrossCoreWaitFlag<4, PIPE_V>(mm2ToReFlag);561 Arch::CrossCoreWaitFlag<4, PIPE_V>(mm2ToReFlag);
601 Arch::CrossCoreSetFlag<4, PIPE_V>(mm2ToReFlag);562 Arch::CrossCoreSetFlag<4, PIPE_V>(mm2ToReFlag);
@@ -614,7 +575,6 @@ private:
614 AscendC::LocalTensor<float> lseUbTensor32;575 AscendC::LocalTensor<float> lseUbTensor32;
615 576 
616 CopyUbToGmO copyUbToGmO;577 CopyUbToGmO copyUbToGmO;
617- 
618};578};
619}579}
620#endif580#endif
Mattention/block_sparse_attention/op_kernel/block_sparse_attention_kernel_interface.cpp+2-2
@@ -185,7 +185,7 @@ __global__ __aicore__ void BsaInferIntfRegular(
185 using TileCopyRescaleO = Epilogue::Tile::TileCopyRescaleO<185 using TileCopyRescaleO = Epilogue::Tile::TileCopyRescaleO<
186 ArchTag, ElementO, LayoutO, LayoutOTmp>;186 ArchTag, ElementO, LayoutO, LayoutOTmp>;
187 using EpilogueRescaleO = Epilogue::Block::BlockEpilogue<187 using EpilogueRescaleO = Epilogue::Block::BlockEpilogue<
188- DispatchPolicyRescaleO, ElementO, ElementOTmp, ElementS, TileCopyRescaleO, Arch::PositionL0C>;188+ DispatchPolicyRescaleO, ElementO, ElementOTmp, ElementS, ElementK, TileCopyRescaleO, Arch::PositionL0C>;
189 189 
190 using BsaRegularKernelArch35 = BsaRegularKernelArch35<190 using BsaRegularKernelArch35 = BsaRegularKernelArch35<
191 EpilogueMask2Idx, BlockMmadQK, EpilogueOnlineSoftmax, BlockMmadPV, EpilogueRescaleO, qFormat, kvFormat>;191 EpilogueMask2Idx, BlockMmadQK, EpilogueOnlineSoftmax, BlockMmadPV, EpilogueRescaleO, qFormat, kvFormat>;
@@ -263,7 +263,7 @@ __global__ __aicore__ void BsaInferInterfaceFullQuant(
263 using DispatchPolicyRescaleO = Epilogue::EpilogueAtlasA5BsaRescaleO<lseMode, lseFormat>;263 using DispatchPolicyRescaleO = Epilogue::EpilogueAtlasA5BsaRescaleO<lseMode, lseFormat>;
264 using TileCopyRescaleO = Epilogue::Tile::TileCopyRescaleO<ArchTag, ElementO, LayoutO, LayoutOTmp>;264 using TileCopyRescaleO = Epilogue::Tile::TileCopyRescaleO<ArchTag, ElementO, LayoutO, LayoutOTmp>;
265 using EpilogueRescaleO = Epilogue::Block::BlockEpilogue<DispatchPolicyRescaleO, ElementO, ElementOTmp, ElementS,265 using EpilogueRescaleO = Epilogue::Block::BlockEpilogue<DispatchPolicyRescaleO, ElementO, ElementOTmp, ElementS,
266- TileCopyRescaleO, Arch::PositionL0C>;266+ ElementK, TileCopyRescaleO, Arch::PositionL0C>;
267 BsaFullQuantKernelParamsArch35 params{query,267 BsaFullQuantKernelParamsArch35 params{query,
268 key,268 key,
269 value,269 value,