Star1137
2751
代码
代码
Issues234
Pull Requests396
Actions
项目讨论
Wiki
分析
Star1137
2751
已合并
smla perf #9988
郑文惠创建于 26 天前
smla perf #9988
已合并
郑文惠创建于 26 天前
从zhengwenhui0817/ops-transformer_2smla_perf合入到cann/ops-transformermaster
共 13 个文件变更+1781-859
Mattention/sparse_flash_mla/op_host/sparse_flash_mla_tiling.cpp+67-6
@@ -353,9 +353,9 @@ ge::graphStatus SMLAInfoParser::GetNpuInfo()
353 if (npuArch_ != NpuArch::DAV_2201 && npuArch_ != NpuArch::DAV_3510) {353 if (npuArch_ != NpuArch::DAV_2201 && npuArch_ != NpuArch::DAV_3510) {
354 OP_LOGE(opName_, "Npu Arch Version[%d] is not support.", (int32_t)npuArch_);354 OP_LOGE(opName_, "Npu Arch Version[%d] is not support.", (int32_t)npuArch_);
355 return ge::GRAPH_FAILED;355 return ge::GRAPH_FAILED;
356- }356+ }
357- batchConsistency_ = false;357+ batchConsistency_ = false;
358- OP_LOGD(opName_, "deterministic_level=%d", context_->GetDeterministicLevel());358+ OP_LOGD(opName_, "deterministic_level=%d", context_->GetDeterministicLevel());
359 359 
360 return ge::GRAPH_SUCCESS;360 return ge::GRAPH_SUCCESS;
361}361}
@@ -889,7 +889,6 @@ ge::graphStatus SMLAInfoParser::GetValueHeadDim()
889 return ge::GRAPH_SUCCESS;889 return ge::GRAPH_SUCCESS;
890}890}
891 891 
892- 
893ge::graphStatus SMLAInfoParser::GetSparseBlockCount()892ge::graphStatus SMLAInfoParser::GetSparseBlockCount()
894{893{
895 if (opParamInfo_.oriSparseIndices.tensor != nullptr) {894 if (opParamInfo_.oriSparseIndices.tensor != nullptr) {
@@ -1990,7 +1989,7 @@ void SparseFlashMlaTiling::CalcUbBmm(SMLATilingInfo *tilingInfo)
1990 cubeMSize = maxMSize;1989 cubeMSize = maxMSize;
1991 }1990 }
1992 mmResUbSize_ = sInnerSizeAlign_ * Align(cubeMSize, 16U); // kernel按照16对齐写出，tiling按照这个原则分配内存1991 mmResUbSize_ = sInnerSizeAlign_ * Align(cubeMSize, 16U); // kernel按照16对齐写出，tiling按照这个原则分配内存
1993- bmm2ResUbSize_ = headDimAlign_ * Align(cubeMSize, 16U); // kernel按照16对齐写出，tiling按照这个原则分配内存1992+ bmm2ResUbSize_ = headDimAlign_ * Align(cubeMSize, 16U); // kernel按照16对齐写出，tiling按照这个原则分配内存
1994}1993}
1995 1994 
1996void SparseFlashMlaTiling::SplitBalanced(SMLATilingInfo *tilingInfo)1995void SparseFlashMlaTiling::SplitBalanced(SMLATilingInfo *tilingInfo)
@@ -2039,6 +2038,63 @@ uint64_t SparseFlashMlaTiling::CalcFdStagingWorkspaceSize(const SMLATilingInfo *
2039 return logicalCoreSlots * FD_MAX_S2_SPLIT_NUM * bytesPerSlot;2038 return logicalCoreSlots * FD_MAX_S2_SPLIT_NUM * bytesPerSlot;
2040}2039}
2041 2040 
2041+uint64_t SparseFlashMlaTiling::CalcVectorizeKvPhyAddrWorkspaceSize(const SMLATilingInfo *tilingInfo,
2042+ uint32_t &vectorizeFlag) const
2043+{
2044+ vectorizeFlag = 0U;
2045+ if (tilingInfo->npuArch != NpuArch::DAV_3510) {
2046+ return 0ULL;
2047+ }
2048+ constexpr uint32_t SPARSE_BLOCK_ALIGN_NUM = 128;
2049+ constexpr uint32_t UB_SIZE = 248 * 1024;
2050+ uint32_t alignedOriSparseBlockCount =
2051+ (tilingInfo->oriSparseBlockCount + SPARSE_BLOCK_ALIGN_NUM - 1) / SPARSE_BLOCK_ALIGN_NUM * SPARSE_BLOCK_ALIGN_NUM;
2052+ uint32_t alignedCmpSparseBlockCount =
2053+ (tilingInfo->cmpSparseBlockCount + SPARSE_BLOCK_ALIGN_NUM - 1) / SPARSE_BLOCK_ALIGN_NUM * SPARSE_BLOCK_ALIGN_NUM;
2054+ bool isPa = (tilingInfo->kvLayout == SMLALayout::PA_BBND);
2055+ uint32_t oriBlocksizeFlag =
2056+ static_cast<uint32_t>((tilingInfo->oriBlockSize & (tilingInfo->oriBlockSize - 1)) == 0);
2057+ uint32_t cmpBlocksizeFlag =
2058+ static_cast<uint32_t>((tilingInfo->cmpBlockSize & (tilingInfo->cmpBlockSize - 1)) == 0);
2059+ uint32_t blocksizeFlag = isPa ? ((tilingInfo->perfMode == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE) ?
2060+ oriBlocksizeFlag :
2061+ (oriBlocksizeFlag && cmpBlocksizeFlag)) :
2062+ 1U;
2063+ uint64_t paExtraUb = std::max(
2064+ static_cast<uint64_t>(tilingInfo->oriMaxBlockNumPerBatch) * sizeof(int32_t),
2065+ static_cast<uint64_t>(tilingInfo->cmpMaxBlockNumPerBatch) * sizeof(int32_t));
2066+ uint64_t cmpUbSize = (isPa ? paExtraUb : 0U) +
2067+ static_cast<uint64_t>(alignedCmpSparseBlockCount) * sizeof(int32_t) +
2068+ static_cast<uint64_t>(alignedCmpSparseBlockCount) * sizeof(int64_t);
2069+ uint64_t oriUbSize = (isPa ? paExtraUb : 0U) +
2070+ static_cast<uint64_t>(alignedOriSparseBlockCount) * sizeof(int32_t) +
2071+ static_cast<uint64_t>(alignedOriSparseBlockCount) * sizeof(int64_t);
2072+ uint64_t vectorizeUbSize = std::max(cmpUbSize, oriUbSize);
2073+ vectorizeFlag = static_cast<uint32_t>(
2074+ (tilingInfo->perfMode == SMLATemplateMode::CSA_TEMPLATE_MODE ||
2075+ tilingInfo->perfMode == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||
2076+ tilingInfo->perfMode == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) &&
2077+ (vectorizeUbSize <= UB_SIZE) && blocksizeFlag);
2078+ 
2079+ if (vectorizeFlag == 0U) {
2080+ return 0ULL;
2081+ }
2082+ uint32_t totalBS1 = (tilingInfo->qLayout == SMLALayout::TND) ?
2083+ tilingInfo->s1Size :
2084+ (tilingInfo->bSize * tilingInfo->s1Size);
2085+ uint64_t oriPhyAddrSize = 0;
2086+ if (tilingInfo->perfMode == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||
2087+ tilingInfo->perfMode == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) {
2088+ oriPhyAddrSize = static_cast<uint64_t>(totalBS1) * alignedOriSparseBlockCount * sizeof(int64_t);
2089+ }
2090+ uint64_t cmpPhyAddrSize = 0;
2091+ if (tilingInfo->perfMode == SMLATemplateMode::CSA_TEMPLATE_MODE ||
2092+ tilingInfo->perfMode == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) {
2093+ cmpPhyAddrSize = static_cast<uint64_t>(totalBS1) * alignedCmpSparseBlockCount * sizeof(int64_t);
2094+ }
2095+ return oriPhyAddrSize + cmpPhyAddrSize;
2096+}
2097+ 
2042// --------------------------SparseFlashMlaTiling类成员函数定义----------------------2098// --------------------------SparseFlashMlaTiling类成员函数定义----------------------
2043ge::graphStatus SparseFlashMlaTiling::DoOpTiling(SMLATilingInfo *tilingInfo)2099ge::graphStatus SparseFlashMlaTiling::DoOpTiling(SMLATilingInfo *tilingInfo)
2044{2100{
@@ -2087,6 +2143,11 @@ ge::graphStatus SparseFlashMlaTiling::DoOpTiling(SMLATilingInfo *tilingInfo)
2087 workspaceSize += MERGE_CACHE_GM_BUF_NUM * 512 * 512 * 2 * aicNum;2143 workspaceSize += MERGE_CACHE_GM_BUF_NUM * 512 * 512 * 2 * aicNum;
2088 }2144 }
2089 }2145 }
2146+ 
2147+ // 计算vectorizeFlag (稀疏KV物理地址向量化)
2148+ uint32_t vectorizeFlag = 0U;
2149+ workspaceSize += CalcVectorizeKvPhyAddrWorkspaceSize(tilingInfo, vectorizeFlag);
2150+ 
2090 workspaceSize += CalcFdStagingWorkspaceSize(tilingInfo, aicNum);2151 workspaceSize += CalcFdStagingWorkspaceSize(tilingInfo, aicNum);
2091 size_t *workSpaces = context_->GetWorkspaceSizes(1);2152 size_t *workSpaces = context_->GetWorkspaceSizes(1);
2092 workSpaces[0] = workspaceSize;2153 workSpaces[0] = workspaceSize;
@@ -2147,7 +2208,7 @@ ge::graphStatus SparseFlashMlaTiling::DoOpTiling(SMLATilingInfo *tilingInfo)
2147 splitG = static_cast<uint32_t>(tilingInfo->gSize > 64);2208 splitG = static_cast<uint32_t>(tilingInfo->gSize > 64);
2148 }2209 }
2149 tilingKey = GET_TPL_TILING_KEY(0U, qLayout, inputKvLayout, static_cast<uint32_t>(tilingInfo->perfMode), splitG,2210 tilingKey = GET_TPL_TILING_KEY(0U, qLayout, inputKvLayout, static_cast<uint32_t>(tilingInfo->perfMode), splitG,
2150- headRatioOne, static_cast<uint32_t>(tilingInfo->batchConsistency));2211+ headRatioOne, static_cast<uint32_t>(tilingInfo->batchConsistency), vectorizeFlag);
2151 context_->SetScheduleMode(1);2212 context_->SetScheduleMode(1);
2152 context_->SetTilingKey(tilingKey);2213 context_->SetTilingKey(tilingKey);
2153 2214 
Mattention/sparse_flash_mla/op_host/sparse_flash_mla_tiling.h+7-3
@@ -313,7 +313,8 @@ public:
313// -----------算子Tiling入参信息解析及Check类---------------313// -----------算子Tiling入参信息解析及Check类---------------
314class SMLATilingCheck {314class SMLATilingCheck {
315public:315public:
316- explicit SMLATilingCheck(const SMLATilingInfo &smlaInfo) : smlaInfo_(smlaInfo) {};316+ explicit SMLATilingCheck(const SMLATilingInfo &smlaInfo)
317+ : smlaInfo_(smlaInfo) {};
317 ~SMLATilingCheck() = default;318 ~SMLATilingCheck() = default;
318 virtual ge::graphStatus Process();319 virtual ge::graphStatus Process();
319 320 
@@ -455,7 +456,8 @@ inline T Align(T num, T rnd)
455 456 
456class SMLAInfoParser {457class SMLAInfoParser {
457public:458public:
458- explicit SMLAInfoParser(gert::TilingContext *context) : context_(context)459+ explicit SMLAInfoParser(gert::TilingContext *context)
460+ : context_(context)
459 {461 {
460 }462 }
461 ~SMLAInfoParser() = default;463 ~SMLAInfoParser() = default;
@@ -589,7 +591,8 @@ public:
589// ---------------算子Tiling类---------------591// ---------------算子Tiling类---------------
590class SparseFlashMlaTiling {592class SparseFlashMlaTiling {
591public:593public:
592- explicit SparseFlashMlaTiling(gert::TilingContext *context) : context_(context) {};594+ explicit SparseFlashMlaTiling(gert::TilingContext *context)
595+ : context_(context) {};
593 ge::graphStatus DoOpTiling(SMLATilingInfo *tilingInfo);596 ge::graphStatus DoOpTiling(SMLATilingInfo *tilingInfo);
594 597 
595private:598private:
@@ -597,6 +600,7 @@ private:
597 void CalcUbBmm(SMLATilingInfo *tilingInfo);600 void CalcUbBmm(SMLATilingInfo *tilingInfo);
598 uint32_t CalcFdLogicalSlotCount(const SMLATilingInfo *tilingInfo, uint32_t aicNum) const;601 uint32_t CalcFdLogicalSlotCount(const SMLATilingInfo *tilingInfo, uint32_t aicNum) const;
599 uint64_t CalcFdStagingWorkspaceSize(const SMLATilingInfo *tilingInfo, uint32_t aicNum) const;602 uint64_t CalcFdStagingWorkspaceSize(const SMLATilingInfo *tilingInfo, uint32_t aicNum) const;
603+ uint64_t CalcVectorizeKvPhyAddrWorkspaceSize(const SMLATilingInfo *tilingInfo, uint32_t &vectorizeFlag) const;
600 gert::TilingContext *context_ = nullptr;604 gert::TilingContext *context_ = nullptr;
601 SMLATemplateMode perfMode_ = SMLATemplateMode::SWA_TEMPLATE_MODE;605 SMLATemplateMode perfMode_ = SMLATemplateMode::SWA_TEMPLATE_MODE;
602 SparseFlashMlaTilingData tilingData_;606 SparseFlashMlaTilingData tilingData_;
Mattention/sparse_flash_mla/op_kernel/arch35/sparse_flash_mla_common_arch35.h+21-14
@@ -21,12 +21,16 @@
21constexpr uint64_t BLOCK_BYTE = 32;21constexpr uint64_t BLOCK_BYTE = 32;
22constexpr uint32_t NEGATIVE_MIN_VAULE_FP32 = 0xFF7FFFFF;22constexpr uint32_t NEGATIVE_MIN_VAULE_FP32 = 0xFF7FFFFF;
23 23 
24-constexpr uint32_t L0AB_SHARED_SIZE_64K = 65536; // 65536表示64*102424+constexpr uint32_t L0AB_SHARED_SIZE_64K = 65536; // 65536表示64*1024
25constexpr uint32_t L0C_SHARED_SIZE_256K = 262144; // 262144表示256 * 102425constexpr uint32_t L0C_SHARED_SIZE_256K = 262144; // 262144表示256 * 1024
26 26 
27-constexpr uint32_t BUFFER_SIZE_16K = 16384; // 16384表示16 * 102427+constexpr uint32_t BUFFER_SIZE_8K = 8192; // 8192表示8 * 1024
28-constexpr uint32_t BUFFER_SIZE_32K = 32768; // 32768表示32 * 102428+constexpr uint32_t BUFFER_SIZE_16K = 16384; // 16384表示16 * 1024
29+constexpr uint32_t BUFFER_SIZE_32K = 32768; // 32768表示32 * 1024
30+constexpr uint32_t BUFFER_SIZE_64K = 65536; // 65536表示64 * 1024
31+constexpr uint32_t BUFFER_SIZE_96K = 98304; // 98304表示96 * 1024
29constexpr uint32_t BUFFER_SIZE_128K = 131072; // 131072表示128 * 102432constexpr uint32_t BUFFER_SIZE_128K = 131072; // 131072表示128 * 1024
33+constexpr uint32_t BUFFER_SIZE_256K = 262144; // 262144表示256 * 1024
30 34 
31constexpr uint32_t CV_RATIO = 2;35constexpr uint32_t CV_RATIO = 2;
32constexpr uint64_t SYNC_MODE = 4;36constexpr uint64_t SYNC_MODE = 4;
@@ -52,21 +56,23 @@ __aicore__ constexpr uint64_t Align64Func(uint64_t data)
52{56{
53 return (data + 63UL) >> 6UL << 6UL; // 向上64对齐, +63移位657 return (data + 63UL) >> 6UL << 6UL; // 向上64对齐, +63移位6
54}58}
55-}59+} // namespace SMLAKernel
56 60 
57#define TEMPLATE_INTF \61#define TEMPLATE_INTF \
58 template <typename Q_T, typename KV_T, typename T, typename OUTPUT_T, \62 template <typename Q_T, typename KV_T, typename T, typename OUTPUT_T, \
59- bool IS_FD, SMLA_LAYOUT LAYOUT_T, \63+ bool IS_FD, SMLA_LAYOUT LAYOUT_T, \
60- SMLA_LAYOUT KV_LAYOUT_T, SMLATemplateMode TEMPLATE_MODE, bool IS_SPLIT_G, bool IS_BATCH_CONSISTENCY>64+ SMLA_LAYOUT KV_LAYOUT_T, SMLATemplateMode TEMPLATE_MODE, bool IS_SPLIT_G, bool IS_BATCH_CONSISTENCY, \
65+ bool IS_VEC_S2PHYADDR>
61 66 
62#define TEMPLATE_INTF_ARGS \67#define TEMPLATE_INTF_ARGS \
63- Q_T, KV_T, T, OUTPUT_T, IS_FD, LAYOUT_T, KV_LAYOUT_T, TEMPLATE_MODE, IS_SPLIT_G, IS_BATCH_CONSISTENCY68+ Q_T, KV_T, T, OUTPUT_T, IS_FD, LAYOUT_T, KV_LAYOUT_T, TEMPLATE_MODE, IS_SPLIT_G, IS_BATCH_CONSISTENCY, \
69+ IS_VEC_S2PHYADDR
64 70 
65#define CUBE_BLOCK_TRAITS_TYPE_FIELDS(X) \71#define CUBE_BLOCK_TRAITS_TYPE_FIELDS(X) \
66 X(Q_T) \72 X(Q_T) \
67 X(KV_T) \73 X(KV_T) \
68 X(T) \74 X(T) \
69- X(OUTPUT_T) \75+ X(OUTPUT_T)
70 76 
71#define CUBE_BLOCK_TRAITS_CONST_FIELDS(X) \77#define CUBE_BLOCK_TRAITS_CONST_FIELDS(X) \
72 X(IS_FD, bool, false) \78 X(IS_FD, bool, false) \
@@ -75,27 +81,28 @@ __aicore__ constexpr uint64_t Align64Func(uint64_t data)
75 X(TEMPLATE_MODE, SMLATemplateMode, SMLATemplateMode::CSA_TEMPLATE_MODE) \81 X(TEMPLATE_MODE, SMLATemplateMode, SMLATemplateMode::CSA_TEMPLATE_MODE) \
76 X(IS_SPLIT_G, bool, false) \82 X(IS_SPLIT_G, bool, false) \
77 X(IS_BATCH_CONSISTENCY, bool, false) \83 X(IS_BATCH_CONSISTENCY, bool, false) \
78- 84+ X(IS_VEC_S2PHYADDR, bool, false)
79 85 
80/* 1. 生成带默认值的模版Template */86/* 1. 生成带默认值的模版Template */
81#define GEN_TYPE_PARAM(name) typename name,87#define GEN_TYPE_PARAM(name) typename name,
82#define GEN_CONST_PARAM(name, type, default_val) type name = default_val,88#define GEN_CONST_PARAM(name, type, default_val) type name = default_val,
83 89 
84#define TEMPLATES_DEF \90#define TEMPLATES_DEF \
85-template <CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_TYPE_PARAM) \91+ template <CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_TYPE_PARAM) \
86- CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_CONST_PARAM) bool end = true>92+ CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_CONST_PARAM) bool end = true>
87 93 
88/* 2. 生成不带带默认值的模版Template */94/* 2. 生成不带带默认值的模版Template */
89#define GEN_TEMPLATE_TYPE_NODEF(name) typename name,95#define GEN_TEMPLATE_TYPE_NODEF(name) typename name,
90#define GEN_TEMPLATE_CONST_NODEF(name, type, default_val) type name,96#define GEN_TEMPLATE_CONST_NODEF(name, type, default_val) type name,
91#define TEMPLATES_DEF_NO_DEFAULT \97#define TEMPLATES_DEF_NO_DEFAULT \
92-template <CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_TEMPLATE_TYPE_NODEF) \98+ template <CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_TEMPLATE_TYPE_NODEF) \
93- CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_TEMPLATE_CONST_NODEF) bool end>99+ CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_TEMPLATE_CONST_NODEF) bool end>
94 100 
95/* 3. 生成有默认值的Args */101/* 3. 生成有默认值的Args */
96#define GEN_ARG_NAME(name, ...) name,102#define GEN_ARG_NAME(name, ...) name,
97#define TEMPLATE_ARGS \103#define TEMPLATE_ARGS \
98 CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_ARG_NAME) \104 CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_ARG_NAME) \
99- CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_ARG_NAME) end105+ CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_ARG_NAME) \
106+ end
100 107 
101#endif108#endif
Mattention/sparse_flash_mla/op_kernel/arch35/sparse_flash_mla_csa_block_cube.h+225-353
@@ -80,10 +80,16 @@ public:
80 static constexpr uint32_t s2BaseSize = 128;80 static constexpr uint32_t s2BaseSize = 128;
81 static constexpr uint32_t dBaseSize = 512;81 static constexpr uint32_t dBaseSize = 512;
82 static constexpr uint32_t dBaseMatmulSize = 128;82 static constexpr uint32_t dBaseMatmulSize = 128;
83+ static constexpr uint32_t rightBufNum = 3;
84+ static constexpr uint32_t rightBufSingleSize = s2BaseSize * dBaseSize;
85+ static constexpr uint32_t rightBufTotalSize = rightBufSingleSize * rightBufNum;
86+ static constexpr uint32_t l1QBufNum = 3; // L1 Q 三缓冲
87+ static constexpr uint32_t l1KBufNum = 3; // L1 K 三缓冲
88+ static constexpr uint32_t qHalfNum = 2; // Q 沿 d 轴切半
89+ static constexpr uint32_t crossCoreMte2SyncFlagId = 15; // IS_SPLIT_G 核间 MTE2 同步 flag ID
83 90 
84 __aicore__ inline CSABlockCube(){};91 __aicore__ inline CSABlockCube(){};
85- __aicore__ inline void InitCubeBlock(TPipe *pipe, BufferManager<BufferType::L1> &l1BufferManager,92+ __aicore__ inline void InitCubeBlock(TPipe *pipe, __gm__ uint8_t *query);
86- __gm__ uint8_t *query);
87 __aicore__ inline void InitCubeInput(__gm__ uint8_t *oriKv, __gm__ uint8_t *cmpKv, __gm__ uint8_t *cmpSparseIndices,93 __aicore__ inline void InitCubeInput(__gm__ uint8_t *oriKv, __gm__ uint8_t *cmpKv, __gm__ uint8_t *cmpSparseIndices,
88 __gm__ uint8_t *oriBlockTable, __gm__ uint8_t *cmpBlockTable,94 __gm__ uint8_t *oriBlockTable, __gm__ uint8_t *cmpBlockTable,
89 __gm__ uint8_t *sequsedQ, __gm__ uint8_t *cuSeqlensQ,95 __gm__ uint8_t *sequsedQ, __gm__ uint8_t *cuSeqlensQ,
@@ -91,52 +97,39 @@ public:
91 __gm__ uint8_t *seqUsedOriKV, __gm__ uint8_t *seqUsedCmpKV,97 __gm__ uint8_t *seqUsedOriKV, __gm__ uint8_t *seqUsedCmpKV,
92 const ConstInfo &constInfo);98 const ConstInfo &constInfo);
93 99 
94- // SWA/HCA场景, inputRightBuf是INNER_CORE_SYNC类型100+ // SWA/HCA场景
101+ __aicore__ inline void IterateLoadQK(RunInfo &runInfo, ConstInfo &constInfo, bool isFirstLoop);
102+ 
103+ // CSA场景
104+ __aicore__ inline void IterateLoadQK(Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,
105+ RunInfo &runInfo, ConstInfo &constInfo, bool isFirstLoop);
95 __aicore__ inline void IterateBmm1(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &output,106 __aicore__ inline void IterateBmm1(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &output,
96- Buffer<BufferType::L1, SyncType::INNER_CORE_SYNC> &inputRightBuf,107+ bool notLastTwoLoop, RunInfo &runInfoNext,
97- const RunInfo &runInfo, const ConstInfo &constInfo);108+ RunInfo &runInfo, ConstInfo &constInfo);
98 __aicore__ inline void109 __aicore__ inline void
99 IterateBmm2(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,110 IterateBmm2(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
100 BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputLeftBuffers,111 BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputLeftBuffers,
101- Buffer<BufferType::L1, SyncType::INNER_CORE_SYNC> &inputRightBuf, const RunInfo &runInfo,112+ const RunInfo &runInfo, const ConstInfo &constInfo);
102- const ConstInfo &constInfo);
103- // CSA场景, inputRightBuf是CROSS_CORE_SYNC_FORWARD类型
104- __aicore__ inline void IterateBmm1(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &output,
105- Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputRightBuf,
106- Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,
107- const RunInfo &runInfo, const ConstInfo &constInfo);
108- __aicore__ inline void
109- IterateBmm2(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
110- BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputLeftBuffers,
111- Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputRightBuf, const RunInfo &runInfo,
112- const ConstInfo &constInfo);
113 113 
114private:114private:
115- __aicore__ inline void InitLocalBuffer(BufferManager<BufferType::L1> &l1BufferManager);115+ __aicore__ inline void InitLocalBuffer();
116 __aicore__ inline void InitGmTensor(__gm__ uint8_t *cuSeqlensQ, __gm__ uint8_t *sequsedQ,116 __aicore__ inline void InitGmTensor(__gm__ uint8_t *cuSeqlensQ, __gm__ uint8_t *sequsedQ,
117 __gm__ uint8_t *cuSeqlensOriKv, __gm__ uint8_t *cuSeqlensCmpKv,117 __gm__ uint8_t *cuSeqlensOriKv, __gm__ uint8_t *cuSeqlensCmpKv,
118 __gm__ uint8_t *seqUsedOriKV, __gm__ uint8_t *seqUsedCmpKV,118 __gm__ uint8_t *seqUsedOriKV, __gm__ uint8_t *seqUsedCmpKV,
119 const ConstInfo &constInfo);119 const ConstInfo &constInfo);
120- __aicore__ inline void CalcS1Coord(const RunInfo &runInfo, const ConstInfo &constInfo);
121 __aicore__ inline void CalcS2Coord(const RunInfo &runInfo, const ConstInfo &constInfo);120 __aicore__ inline void CalcS2Coord(const RunInfo &runInfo, const ConstInfo &constInfo);
122- __aicore__ inline void IterateBmm1HCA(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,121+ __aicore__ inline void CopyQGmToL1(RunInfo &runInfo, ConstInfo &constInfo);
123- Buffer<BufferType::L1, SyncType::INNER_CORE_SYNC> &inputRightBuf,122+ __aicore__ inline void LoadKGmToL1(LocalTensor<KV_T> &inputRightTensor, const RunInfo &runInfo,
124- const RunInfo &runInfo, const ConstInfo &constInfo);123+ const ConstInfo &constInfo);
125- __aicore__ inline void
126- IterateBmm2HCA(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
127- BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputLeftBuffers,
128- Buffer<BufferType::L1, SyncType::INNER_CORE_SYNC> &inputRightBuf, const RunInfo &runInfo,
129- const ConstInfo &constInfo);
130 124 
131- __aicore__ inline void IterateBmm1CSA(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,125+ __aicore__ inline void IterateBmm1Impl(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
132- Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputRightBuf,126+ bool notLastTwoLoop,
133- Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,127+ RunInfo &runInfoNext, RunInfo &runInfo, ConstInfo &constInfo);
134- const RunInfo &runInfo, const ConstInfo &constInfo);
135 __aicore__ inline void128 __aicore__ inline void
136- IterateBmm2CSA(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,129+ IterateBmm2Impl(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
137- BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputLeftBuffers,130+ BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputLeftBuffers,
138- Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputRightBuf, const RunInfo &runInfo,131+ const RunInfo &runInfo, const ConstInfo &constInfo);
139- const ConstInfo &constInfo);132+ 
140 TPipe *tPipe;133 TPipe *tPipe;
141 /* =====================GM变量==================== */134 /* =====================GM变量==================== */
142 static constexpr GmFormat Q_FORMAT = GetQueryGmFormat<LAYOUT_T>();135 static constexpr GmFormat Q_FORMAT = GetQueryGmFormat<LAYOUT_T>();
@@ -157,35 +150,37 @@ private:
157 CubeCoordInfo coordInfo[3];150 CubeCoordInfo coordInfo[3];
158 uint32_t kvCacheBlockSize = 0;151 uint32_t kvCacheBlockSize = 0;
159 uint32_t maxBlockNumPerBatch = 0;152 uint32_t maxBlockNumPerBatch = 0;
160- TEventID mte1ToMte2Id[3];153+ uint32_t l0CBufId = 0;
161- TEventID mte2ToMte1Id[3];154+ uint32_t l1QBufId = 0;
162- 155+ uint32_t l1KLoadBufId = 0;
156+ uint32_t l1KMatmul1BufId = 0;
157+ uint32_t l1KMatmul2BufId = 0;
158+ uint32_t l0CFixToMFlagId = 0; // {0, 1}, 用于L0C
159+ uint32_t l0CMToFixFlagId = 0; // {0, 1}, 用于L0C
160+ uint32_t l1QMte1ToMte2FlagId = 0; // {0, 1, 2}, 用于l1Q
161+ uint32_t l1QMte2ToMte1FlagId = 0; // {0, 1, 2}, 用于l1Q
162+ uint32_t l1KMte1ToMte2FlagId = 3; // {3, 4, 5}, 用于l1K
163+ uint32_t l1KMte2ToMte1FlagId = 3; // {3, 4, 5}, 用于l1K
163 /* =====================LocalBuffer变量==================== */164 /* =====================LocalBuffer变量==================== */
164 BufferManager<BufferType::L0A> l0aBufferManager;165 BufferManager<BufferType::L0A> l0aBufferManager;
165 BufferManager<BufferType::L0B> l0bBufferManager;166 BufferManager<BufferType::L0B> l0bBufferManager;
166- BufferManager<BufferType::L0C> l0cBufferManager;
167- 
168- // D小于等于256 mm1左矩阵Q，GS1循环内左矩阵复用,
169- // GS1循环间开pingpong；D大于256使用单块Buffer，S1循环间驻留；fp32场景单块不驻留
170- BuffersPolicySingleBuffer<BufferType::L1> l1QBuffers;
171- 
172- // L0A
173 BuffersPolicyDB<BufferType::L0A> mmL0ABuffers;167 BuffersPolicyDB<BufferType::L0A> mmL0ABuffers;
174- // L0B
175 BuffersPolicyDB<BufferType::L0B> mmL0BBuffers;168 BuffersPolicyDB<BufferType::L0B> mmL0BBuffers;
176- // L0C169+ TBuf<TPosition::B1> l1RightBuffers;
177- BuffersPolicyDB<BufferType::L0C> mmL0CBuffers;170+ TBuf<TPosition::A1> l1QBuffers;
171+ TBuf<TPosition::CO1> mmL0CBuffers;
172+ LocalTensor<Q_T> l1RightTensor;
173+ LocalTensor<Q_T> l1QTensor;
174+ LocalTensor<T> mmL0CTensor;
178};175};
179 176 
180TEMPLATES_DEF_NO_DEFAULT177TEMPLATES_DEF_NO_DEFAULT
181-__aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::InitCubeBlock(TPipe *pipe,178+__aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::InitCubeBlock(TPipe *pipe, __gm__ uint8_t *query)
182- BufferManager<BufferType::L1> &l1BufferManager,
183- __gm__ uint8_t *query)
184{179{
185 if ASCEND_IS_AIC {180 if ASCEND_IS_AIC {
186 tPipe = pipe;181 tPipe = pipe;
187 this->queryGm.gmTensor.SetGlobalBuffer((__gm__ Q_T *)query);182 this->queryGm.gmTensor.SetGlobalBuffer((__gm__ Q_T *)query);
188- InitLocalBuffer(l1BufferManager);183+ InitLocalBuffer();
189 }184 }
190}185}
191 186 
@@ -211,37 +206,35 @@ __aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::InitCubeInput(
211 if constexpr (TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE) {206 if constexpr (TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE) {
212 this->cmpSparseIndicesGm.SetGlobalBuffer((__gm__ int32_t *)cmpSparseIndices);207 this->cmpSparseIndicesGm.SetGlobalBuffer((__gm__ int32_t *)cmpSparseIndices);
213 }208 }
214- if constexpr (TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE ||
215- TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||
216- TEMPLATE_MODE == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE || IS_SPLIT_G) {
217- mte1ToMte2Id[0] = GetTPipePtr()->AllocEventID<HardEvent::MTE2_MTE1>();
218- mte1ToMte2Id[1] = GetTPipePtr()->AllocEventID<HardEvent::MTE2_MTE1>();
219- mte1ToMte2Id[2] = GetTPipePtr()->AllocEventID<HardEvent::MTE2_MTE1>();
220- mte2ToMte1Id[0] = GetTPipePtr()->AllocEventID<HardEvent::MTE1_MTE2>();
221- mte2ToMte1Id[1] = GetTPipePtr()->AllocEventID<HardEvent::MTE1_MTE2>();
222- mte2ToMte1Id[2] = GetTPipePtr()->AllocEventID<HardEvent::MTE1_MTE2>();
223- SetFlag<HardEvent::MTE1_MTE2>(mte2ToMte1Id[0]);
224- SetFlag<HardEvent::MTE1_MTE2>(mte2ToMte1Id[1]);
225- SetFlag<HardEvent::MTE1_MTE2>(mte2ToMte1Id[2]);
226- }
227 InitGmTensor(cuSeqlensQ, sequsedQ, cuSeqlensOriKv, cuSeqlensCmpKv, seqUsedOriKV, seqUsedCmpKV, constInfo);209 InitGmTensor(cuSeqlensQ, sequsedQ, cuSeqlensOriKv, cuSeqlensCmpKv, seqUsedOriKV, seqUsedCmpKV, constInfo);
228 }210 }
229}211}
230 212 
231TEMPLATES_DEF_NO_DEFAULT213TEMPLATES_DEF_NO_DEFAULT
232-__aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::InitLocalBuffer(BufferManager<BufferType::L1> &l1BufferManager)214+__aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::InitLocalBuffer()
233{215{
234- constexpr uint32_t mm1LeftSize = s1BaseSize * dBaseSize * sizeof(Q_T);216+ tPipe->InitBuffer(l1RightBuffers, rightBufTotalSize * sizeof(Q_T));
235- l1QBuffers.Init((l1BufferManager), mm1LeftSize);217+ l1RightTensor = l1RightBuffers.Get<Q_T>();
218+ 
219+ tPipe->InitBuffer(l1QBuffers, BUFFER_SIZE_96K);
220+ l1QTensor = l1QBuffers.Get<Q_T>();
236 221 
237- // L0A B C 当前写死，能否通过基础api获取
238 l0aBufferManager.Init(tPipe, L0AB_SHARED_SIZE_64K);222 l0aBufferManager.Init(tPipe, L0AB_SHARED_SIZE_64K);
239 l0bBufferManager.Init(tPipe, L0AB_SHARED_SIZE_64K);223 l0bBufferManager.Init(tPipe, L0AB_SHARED_SIZE_64K);
240- l0cBufferManager.Init(tPipe, L0C_SHARED_SIZE_256K);
241 224 
242- mmL0ABuffers.Init(l0aBufferManager, BUFFER_SIZE_16K); // db类型，填入数值是总大小的一半225+ mmL0ABuffers.Init(l0aBufferManager, BUFFER_SIZE_16K);
243 mmL0BBuffers.Init(l0bBufferManager, BUFFER_SIZE_32K);226 mmL0BBuffers.Init(l0bBufferManager, BUFFER_SIZE_32K);
244- mmL0CBuffers.Init(l0cBufferManager, BUFFER_SIZE_128K);227+ 
228+ tPipe->InitBuffer(mmL0CBuffers, BUFFER_SIZE_256K);
229+ mmL0CTensor = mmL0CBuffers.Get<T>();
230+ SetFlag<HardEvent::FIX_M>(l0CFixToMFlagId); // {0, 1}, 用于L0C
231+ SetFlag<HardEvent::FIX_M>(l0CFixToMFlagId + 1);
232+ SetFlag<HardEvent::MTE1_MTE2>(l1QMte1ToMte2FlagId); // {0, 1, 2}, 用于l1Q
233+ SetFlag<HardEvent::MTE1_MTE2>(l1QMte1ToMte2FlagId + 1);
234+ SetFlag<HardEvent::MTE1_MTE2>(l1QMte1ToMte2FlagId + 2);
235+ SetFlag<HardEvent::MTE1_MTE2>(l1KMte1ToMte2FlagId); // {3, 4, 5}, 用于l1K
236+ SetFlag<HardEvent::MTE1_MTE2>(l1KMte1ToMte2FlagId + 1);
237+ SetFlag<HardEvent::MTE1_MTE2>(l1KMte1ToMte2FlagId + 2);
245}238}
246 239 
247/* 初始化GmTensor,设置shape信息并计算strides */240/* 初始化GmTensor,设置shape信息并计算strides */
@@ -293,13 +286,6 @@ __aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::InitGmTensor(__gm__ uint8_t
293 }286 }
294}287}
295 288 
296-TEMPLATES_DEF_NO_DEFAULT
297-__aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::CalcS1Coord(const RunInfo &runInfo, const ConstInfo &constInfo)
298-{
299- // 计算s1方向偏移
300- coordInfo[runInfo.taskIdMod3].s1Coord = runInfo.s1oIdx * runInfo.qSNumInOneBlock;
301-}
302- 
303TEMPLATES_DEF_NO_DEFAULT289TEMPLATES_DEF_NO_DEFAULT
304__aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::CalcS2Coord(const RunInfo &runInfo, const ConstInfo &constInfo)290__aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::CalcS2Coord(const RunInfo &runInfo, const ConstInfo &constInfo)
305{291{
@@ -322,57 +308,25 @@ __aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::CalcS2Coord(const RunInfo &r
322}308}
323 309 
324TEMPLATES_DEF_NO_DEFAULT310TEMPLATES_DEF_NO_DEFAULT
325-__aicore__ inline void311+__aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::CopyQGmToL1(RunInfo &runInfo, ConstInfo &constInfo)
326-CSABlockCube<TEMPLATE_ARGS>::IterateBmm1(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
327- Buffer<BufferType::L1, SyncType::INNER_CORE_SYNC> &inputRightBuf,
328- const RunInfo &runInfo, const ConstInfo &constInfo)
329{312{
330- CalcS1Coord(runInfo, constInfo);313+ uint64_t gmOffset = this->queryGm.offsetCalculator.GetOffset(runInfo.boIdx, runInfo.n2oIdx, runInfo.goIdx,
331- CalcS2Coord(runInfo, constInfo);314+ runInfo.s1oIdx * runInfo.qSNumInOneBlock, 0);
332- 315+ for (uint32_t i = 0; i < qHalfNum; i++) {
333- IterateBmm1HCA(outputBuf, inputRightBuf, runInfo, constInfo);316+ uint32_t curL1QBufId = (l1QBufId + i) % l1QBufNum;
334-}317+ WaitFlag<HardEvent::MTE1_MTE2>(l1QMte1ToMte2FlagId + curL1QBufId);
335- 318+ uint64_t curGmOffset = gmOffset + i * (constInfo.dSize >> 1);
336-TEMPLATES_DEF_NO_DEFAULT319+ CopyToL1Nd2Nz<Q_T>(l1QTensor[curL1QBufId * BUFFER_SIZE_16K],
337-__aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::IterateBmm2(320+ this->queryGm.gmTensor[curGmOffset], runInfo.mRealSize, constInfo.dSize >> 1,
338- Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
339- BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputLeftBuffers,
340- Buffer<BufferType::L1, SyncType::INNER_CORE_SYNC> &inputRightBuf, const RunInfo &runInfo,
341- const ConstInfo &constInfo)
342-{
343- IterateBmm2HCA(outputBuf, inputLeftBuffers, inputRightBuf, runInfo, constInfo);
344-}
345- 
346-TEMPLATES_DEF_NO_DEFAULT
347-__aicore__ inline void
348-CSABlockCube<TEMPLATE_ARGS>::IterateBmm1HCA(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
349- Buffer<BufferType::L1, SyncType::INNER_CORE_SYNC> &inputRightBuf,
350- const RunInfo &runInfo, const ConstInfo &constInfo)
351-{
352- Buffer<BufferType::L1> inputLeftBuf;
353- // 左矩阵复用，S2的第一次循环加载左矩阵
354- // 加载左矩阵到L1, 全载
355- // query对ori_kv, cmp_kv都一样，无需区分
356- if (unlikely(runInfo.s2LoopCount == 0)) { // sOuter循环第一个基本块：搬运Q
357- inputLeftBuf = l1QBuffers.Get();
358- inputLeftBuf.Wait<HardEvent::MTE1_MTE2>(); // 占用L1A
359- LocalTensor<Q_T> inputLeftTensor = inputLeftBuf.GetTensor<Q_T>();
360- 
361- uint64_t gmOffset = this->queryGm.offsetCalculator.GetOffset(runInfo.boIdx, runInfo.n2oIdx, runInfo.goIdx,
362- coordInfo[runInfo.taskIdMod3].s1Coord, 0);
363- CopyToL1Nd2Nz<Q_T>(inputLeftTensor, this->queryGm.gmTensor[gmOffset], runInfo.mRealSize, constInfo.dSize,
364 constInfo.mm1Ka);321 constInfo.mm1Ka);
365- 322+ SetFlag<HardEvent::MTE2_MTE1>(l1QMte2ToMte1FlagId + curL1QBufId);
366- inputLeftBuf.Set<HardEvent::MTE2_MTE1>(); // 通知
367- } else { // 非S2的第一次循环直接复用Q
368- inputLeftBuf = l1QBuffers.GetPre();
369- // 左矩阵复用时，sinner循环内不需要MTE2同步等待
370- inputLeftBuf.Set<HardEvent::MTE2_MTE1>(); // 通知
371 }323 }
324+}
372 325 
373- // 加载当前轮的右矩阵到L1326+TEMPLATES_DEF_NO_DEFAULT
374- inputRightBuf.Wait<HardEvent::MTE1_MTE2>(); // 占用L1B327+__aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::LoadKGmToL1(LocalTensor<KV_T> &inputRightTensor,
375- LocalTensor<KV_T> inputRightTensor = inputRightBuf.GetTensor<KV_T>();328+ const RunInfo &runInfo, const ConstInfo &constInfo)
329+{
376 if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND) {330 if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND) {
377 Position startPos;331 Position startPos;
378 startPos.bIdx = runInfo.boIdx;332 startPos.bIdx = runInfo.boIdx;
@@ -395,156 +349,36 @@ CSABlockCube<TEMPLATE_ARGS>::IterateBmm1HCA(Buffer<BufferType::UB, SyncType::CRO
395 CopyToL1Nd2Nz<KV_T>(inputRightTensor, curKvGm.gmTensor[keyOffset], runInfo.s2RealSize, constInfo.dSize,349 CopyToL1Nd2Nz<KV_T>(inputRightTensor, curKvGm.gmTensor[keyOffset], runInfo.s2RealSize, constInfo.dSize,
396 constInfo.mm1Kb);350 constInfo.mm1Kb);
397 }351 }
398- inputRightBuf.Set<HardEvent::MTE2_MTE1>(); // 通知
399- inputRightBuf.Wait<HardEvent::MTE2_MTE1>(); // 等待L1B
400- 
401- inputLeftBuf.Wait<HardEvent::MTE2_MTE1>(); // 等待L1A
402- Buffer<BufferType::L0C> mm1ResL0C = mmL0CBuffers.Get();
403- mm1ResL0C.Wait<HardEvent::FIX_M>(); // 占用
404- MMParam param = {
405- static_cast<uint32_t>(runInfo.mRealSize), // singleM
406- static_cast<uint32_t>(runInfo.s2RealSize), // singleN
407- static_cast<uint32_t>(constInfo.dSize), // singleK
408- 0, // isLeftTranspose
409- 1 // isRightTranspose
410- };
411- MatmulK<Q_T, Q_T, T, s1BaseSize, s2BaseSize, dBaseMatmulSize, ABLayout::MK, ABLayout::KN>( // m,n不切，k切128
412- inputLeftBuf.GetTensor<Q_T>(), inputRightBuf.GetTensor<Q_T>(), // mm1B直接用tensor的数据
413- mmL0ABuffers, mmL0BBuffers, mm1ResL0C.GetTensor<T>(), param);
414- if (unlikely(runInfo.s2LoopCount == runInfo.s2LoopLimit)) {
415- inputLeftBuf.Set<HardEvent::MTE1_MTE2>(); // 释放L1A
416- }
417- 
418- mm1ResL0C.Set<HardEvent::M_FIX>(); // 通知
419- mm1ResL0C.Wait<HardEvent::M_FIX>(); // 等待L0C
420- 
421- outputBuf.WaitCrossCore();
422- FixpipeParamsC310<CO2Layout::ROW_MAJOR> fixpipeParams; // L0C→UB
423- // L0C上的bmm1结果矩阵N方向的size大小; 同mmadParams.n; 为什么要8个元素对齐(32B对齐) // 128
424- fixpipeParams.nSize = Align8Func(runInfo.s2RealSize);
425- // 有效数据不足16行，只需要输出部分行即可; L0C上的bmm1结果矩阵M方向的size大小(必须为偶数) // 128
426- fixpipeParams.mSize = Align2Func(runInfo.mRealSize);
427- // L0C上bmm1结果相邻连续数据片段间隔(前面一个数据块的头与后面数据块的头的间隔), 单位为16*sizeof(T)
428- // 源Nz矩阵中相邻大Z排布的起始地址偏移
429- fixpipeParams.srcStride = Align16Func(fixpipeParams.mSize);
430- // mmResUb上两行之间的间隔，单位：element。 // 128:根据比对dump文件得到, ND方案(S1*S2)时脏数据用mask剔除
431- fixpipeParams.dstStride = s2BaseSize;
432- // 双目标模式，按M维度拆分，M / 2 * N写入每个UB, M必须为2的倍数
433- fixpipeParams.dualDstCtl = 1;
434- fixpipeParams.params.ndNum = 1;
435- fixpipeParams.params.srcNdStride = 0;
436- fixpipeParams.params.dstNdStride = 0;
437- 
438- // 将matmul结果从L0C搬运到UB
439- Fixpipe<T, T, PFA_CFG_ROW_MAJOR_UB>(outputBuf.template GetTensor<T>(), mm1ResL0C.GetTensor<T>(), fixpipeParams);
440- mm1ResL0C.Set<HardEvent::FIX_M>(); // 释放L0C
441- outputBuf.SetCrossCore();
442-}
443- 
444-TEMPLATES_DEF_NO_DEFAULT
445-__aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::IterateBmm2HCA(
446- Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
447- BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputLeftBuffers,
448- Buffer<BufferType::L1, SyncType::INNER_CORE_SYNC> &inputRightBuf, const RunInfo &runInfo,
449- const ConstInfo &constInfo)
450-{
451- Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> inputLeftBuf = inputLeftBuffers.Get(); // P直接用无需搬运
452- inputLeftBuf.WaitCrossCore();
453- 
454- Buffer<BufferType::L0C> mm2ResL0C = mmL0CBuffers.Get();
455- mm2ResL0C.Wait<HardEvent::FIX_M>(); // 占用
456- MMParam param = {
457- static_cast<uint32_t>(runInfo.mRealSize), // singleM 64
458- static_cast<uint32_t>(constInfo.dSizeV), // singleN 512
459- static_cast<uint32_t>(runInfo.s2RealSize), // singleK 128
460- 0, // isLeftTranspose
461- 0 // isRightTranspose
462- };
463- MatmulN<Q_T, Q_T, T, s1BaseSize, s2BaseSize, dBaseMatmulSize, ABLayout::MK, ABLayout::KN>(
464- inputLeftBuf.GetTensor<Q_T>(), inputRightBuf.GetTensor<Q_T>(), mmL0ABuffers, mmL0BBuffers,
465- mm2ResL0C.GetTensor<T>(), param);
466- 
467- inputLeftBuf.SetCrossCore();
468- // bmm2才释放KV，在这里释放
469- inputRightBuf.Set<HardEvent::MTE1_MTE2>();
470- 
471- mm2ResL0C.Set<HardEvent::M_FIX>(); // 通知
472- mm2ResL0C.Wait<HardEvent::M_FIX>(); // 等待
473- 
474- // 占用
475- outputBuf.WaitCrossCore();
476- // L0C→UB;FixpipeParamsM300:L0C→UB
477- FixpipeParamsC310<CO2Layout::ROW_MAJOR> fixpipeParams;
478- // L0C上的bmm1结果矩阵N方向的size大小, 分档计算且vector2中通过mask筛选出实际有效值
479- fixpipeParams.nSize = Align8Func(constInfo.dSizeV);
480- // 有效数据不足16行，只需要输出部分行即可; L0C上的bmm1结果矩阵M方向的size大小; 同mmadParams.m
481- fixpipeParams.mSize = Align2Func(runInfo.mRealSize);
482- // L0C上bmm1结果相邻连续数据片段间隔（前面一个数据块的头与后面数据块的头的间隔）
483- fixpipeParams.srcStride = Align16Func(fixpipeParams.mSize);
484- fixpipeParams.dstStride = Align16Func(constInfo.dSizeV);
485- fixpipeParams.dualDstCtl = 1;
486- fixpipeParams.params.ndNum = 1;
487- fixpipeParams.params.srcNdStride = 0;
488- fixpipeParams.params.dstNdStride = 0;
489- Fixpipe<T, T, PFA_CFG_ROW_MAJOR_UB>(outputBuf.template GetTensor<T>(), mm2ResL0C.GetTensor<T>(),
490- fixpipeParams); // 将matmul结果从L0C搬运到UB
491- mm2ResL0C.Set<HardEvent::FIX_M>(); // 释放
492- 
493- outputBuf.SetCrossCore();
494}352}
495 353 
354+// SWA/HCA场景: K从GM直接搬运
496TEMPLATES_DEF_NO_DEFAULT355TEMPLATES_DEF_NO_DEFAULT
497__aicore__ inline void356__aicore__ inline void
498-CSABlockCube<TEMPLATE_ARGS>::IterateBmm1(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,357+CSABlockCube<TEMPLATE_ARGS>::IterateLoadQK(RunInfo &runInfo, ConstInfo &constInfo, bool isFirstLoop)
499- Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputRightBuf,
500- Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,
501- const RunInfo &runInfo, const ConstInfo &constInfo)
502{358{
503- CalcS1Coord(runInfo, constInfo);359+ if (unlikely(isFirstLoop)) {
504- CalcS2Coord(runInfo, constInfo);360+ CopyQGmToL1(runInfo, constInfo);
505- 
506- IterateBmm1CSA(outputBuf, inputRightBuf, v0ResGm, runInfo, constInfo);
507-}
508- 
509-TEMPLATES_DEF_NO_DEFAULT
510-__aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::IterateBmm2(
511- Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
512- BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputLeftBuffers,
513- Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputRightBuf, const RunInfo &runInfo,
514- const ConstInfo &constInfo)
515-{
516- IterateBmm2CSA(outputBuf, inputLeftBuffers, inputRightBuf, runInfo, constInfo);
517-}
518- 
519-TEMPLATES_DEF_NO_DEFAULT
520-__aicore__ inline void
521-CSABlockCube<TEMPLATE_ARGS>::IterateBmm1CSA(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
522- Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputRightBuf,
523- Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,
524- const RunInfo &runInfo, const ConstInfo &constInfo)
525-{
526- Buffer<BufferType::L1> inputLeftBuf;
527- // 左矩阵复用，S2的第一次循环加载左矩阵
528- // 加载左矩阵到L1, 全载
529- // query对ori_kv, cmp_kv都一样，无需区分
530- if (unlikely(runInfo.s2LoopCount == 0)) { // sOuter循环第一个基本块：搬运Q
531- inputLeftBuf = l1QBuffers.Get();
532- inputLeftBuf.Wait<HardEvent::MTE1_MTE2>(); // 占用L1A
533- LocalTensor<Q_T> inputLeftTensor = inputLeftBuf.GetTensor<Q_T>();
534- 
535- uint64_t gmOffset = this->queryGm.offsetCalculator.GetOffset(runInfo.boIdx, runInfo.n2oIdx, runInfo.goIdx,
536- coordInfo[runInfo.taskIdMod3].s1Coord, 0);
537- CopyToL1Nd2Nz<Q_T>(inputLeftTensor, this->queryGm.gmTensor[gmOffset], runInfo.mRealSize, constInfo.dSize,
538- constInfo.mm1Ka);
539- 
540- inputLeftBuf.Set<HardEvent::MTE2_MTE1>(); // 通知
541- } else { // 非S2的第一次循环直接复用Q
542- inputLeftBuf = l1QBuffers.GetPre();
543- // 左矩阵复用时，sinner循环内不需要MTE2同步等待
544- inputLeftBuf.Set<HardEvent::MTE2_MTE1>(); // 通知
545 }361 }
546 362 
547 // 加载当前轮的右矩阵到L1363 // 加载当前轮的右矩阵到L1
364+ CalcS2Coord(runInfo, constInfo);
365+ WaitFlag<HardEvent::MTE1_MTE2>(l1KMte1ToMte2FlagId + l1KLoadBufId);
366+ LocalTensor<KV_T> dst = l1RightTensor[runInfo.taskIdMod3 * rightBufSingleSize];
367+ LoadKGmToL1(dst, runInfo, constInfo);
368+ SetFlag<HardEvent::MTE2_MTE1>(l1KMte2ToMte1FlagId + l1KLoadBufId);
369+ l1KLoadBufId = (l1KLoadBufId + 1) % l1KBufNum;
370+}
371+ 
372+// CSA场景: K来源取决于isSparse
373+TEMPLATES_DEF_NO_DEFAULT
374+__aicore__ inline void
375+CSABlockCube<TEMPLATE_ARGS>::IterateLoadQK(Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,
376+ RunInfo &runInfo, ConstInfo &constInfo, bool isFirstLoop)
377+{
378+ if (unlikely(isFirstLoop)) {
379+ CopyQGmToL1(runInfo, constInfo);
380+ }
381+ 
548 // ORI_SPARSE、ORI_CMP_SPARSE及CSA的cmpKv为v0稀疏搬运，CSA的oriKv为cube搬运382 // ORI_SPARSE、ORI_CMP_SPARSE及CSA的cmpKv为v0稀疏搬运，CSA的oriKv为cube搬运
549 bool isSparse = false;383 bool isSparse = false;
550 if constexpr (TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||384 if constexpr (TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||
@@ -554,67 +388,107 @@ CSABlockCube<TEMPLATE_ARGS>::IterateBmm1CSA(Buffer<BufferType::UB, SyncType::CRO
554 isSparse = runInfo.isCmp ? true : false;388 isSparse = runInfo.isCmp ? true : false;
555 }389 }
556 390 
391+ WaitFlag<HardEvent::MTE1_MTE2>(l1KMte1ToMte2FlagId + l1KLoadBufId);
392+ LocalTensor<Q_T> dst = l1RightTensor[runInfo.taskIdMod3 * rightBufSingleSize];
557 if (!isSparse) {393 if (!isSparse) {
558- WaitFlag<HardEvent::MTE1_MTE2>(mte2ToMte1Id[runInfo.taskIdMod3]);394+ if constexpr (IS_SPLIT_G) {
559- LocalTensor<KV_T> inputRightTensor = inputRightBuf.GetTensor<KV_T>();395+ CrossCoreSetFlag<0, PIPE_MTE2>(crossCoreMte2SyncFlagId);
560- if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND) {396+ CrossCoreWaitFlag<0, PIPE_MTE2>(crossCoreMte2SyncFlagId);
561- Position startPos;
562- startPos.bIdx = runInfo.boIdx;
563- startPos.n2Idx = runInfo.n2oIdx;
564- startPos.s2Offset = coordInfo[runInfo.taskIdMod3].s2Coord;
565- startPos.dIdx = 0;
566- PAShape shape;
567- shape.blockSize = kvCacheBlockSize;
568- shape.headNum = constInfo.n2Size;
569- shape.headDim = constInfo.dSize;
570- shape.actHeadDim = constInfo.dSize;
571- shape.maxblockNumPerBatch = maxBlockNumPerBatch;
572- shape.copyRowNum = runInfo.s2RealSize;
573- shape.copyRowNumAlign = (runInfo.s2RealSize + 15) >> 4 << 4;
574- shape.pageStride = runInfo.isCmp ? constInfo.cmpKeyStride0 : constInfo.oriKeyStride0;
575- GmCopyInToL1PA<KV_T>(inputRightTensor, curKvGm.gmTensor, blockTableGm, KVLAYOUT::BBH, shape, startPos);
576- } else {
577- int64_t keyOffset = this->curKvGm.offsetCalculator.GetOffset(
578- coordInfo[runInfo.taskIdMod3].curBIdx, runInfo.n2oIdx, coordInfo[runInfo.taskIdMod3].s2Coord, 0);
579- CopyToL1Nd2Nz<KV_T>(inputRightTensor, curKvGm.gmTensor[keyOffset], runInfo.s2RealSize, constInfo.dSize,
580- constInfo.mm1Kb);
581 }397 }
582- SetFlag<HardEvent::MTE2_MTE1>(mte1ToMte2Id[runInfo.taskIdMod3]);398+ // cube直接从kv cache搬运K
583- WaitFlag<HardEvent::MTE2_MTE1>(mte1ToMte2Id[runInfo.taskIdMod3]);399+ CalcS2Coord(runInfo, constInfo);
400+ LoadKGmToL1(dst, runInfo, constInfo);
584 } else {401 } else {
585- inputRightBuf.WaitCrossCore();
586- WaitFlag<HardEvent::MTE1_MTE2>(mte2ToMte1Id[runInfo.taskIdMod3]);
587- LocalTensor<Q_T> dst = inputRightBuf.GetTensor<Q_T>();
588 v0ResGm.WaitCrossCore();402 v0ResGm.WaitCrossCore();
403+ if constexpr (IS_SPLIT_G) {
404+ CrossCoreSetFlag<0, PIPE_MTE2>(crossCoreMte2SyncFlagId);
405+ CrossCoreWaitFlag<0, PIPE_MTE2>(crossCoreMte2SyncFlagId);
406+ }
589 GlobalTensor<Q_T> v0ResGmTensor = v0ResGm.template GetTensor<Q_T>();407 GlobalTensor<Q_T> v0ResGmTensor = v0ResGm.template GetTensor<Q_T>();
590 CopyToL1Nd2Nz<Q_T>(dst, v0ResGmTensor, runInfo.s2RealSize, constInfo.dSize, constInfo.mm1Kb);408 CopyToL1Nd2Nz<Q_T>(dst, v0ResGmTensor, runInfo.s2RealSize, constInfo.dSize, constInfo.mm1Kb);
591- SetFlag<HardEvent::MTE2_MTE1>(mte1ToMte2Id[runInfo.taskIdMod3]);
592- WaitFlag<HardEvent::MTE2_MTE1>(mte1ToMte2Id[runInfo.taskIdMod3]);
593 }409 }
410+ SetFlag<HardEvent::MTE2_MTE1>(l1KMte2ToMte1FlagId + l1KLoadBufId);
411+ l1KLoadBufId = (l1KLoadBufId + 1) % l1KBufNum;
412+}
413+ 
414+TEMPLATES_DEF_NO_DEFAULT
415+__aicore__ inline void
416+CSABlockCube<TEMPLATE_ARGS>::IterateBmm1(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
417+ bool notLastTwoLoop, RunInfo &runInfoNext,
418+ RunInfo &runInfo, ConstInfo &constInfo)
419+{
420+ IterateBmm1Impl(outputBuf, notLastTwoLoop, runInfoNext, runInfo, constInfo);
421+}
422+ 
423+TEMPLATES_DEF_NO_DEFAULT
424+__aicore__ inline void
425+CSABlockCube<TEMPLATE_ARGS>::IterateBmm2(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
426+ BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputLeftBuffers,
427+ const RunInfo &runInfo, const ConstInfo &constInfo)
428+{
429+ IterateBmm2Impl(outputBuf, inputLeftBuffers, runInfo, constInfo);
430+}
431+ 
432+TEMPLATES_DEF_NO_DEFAULT
433+__aicore__ inline void
434+CSABlockCube<TEMPLATE_ARGS>::IterateBmm1Impl(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
435+ bool notLastTwoLoop,
436+ RunInfo &runInfoNext, RunInfo &runInfo, ConstInfo &constInfo)
437+{
438+ LocalTensor<Q_T> curL1RightTensor = l1RightTensor[runInfo.taskIdMod3 * rightBufSingleSize];
439+ WaitFlag<HardEvent::MTE2_MTE1>(l1KMte2ToMte1FlagId + l1KMatmul1BufId);
440+ l1KMatmul1BufId = (l1KMatmul1BufId + 1) % l1KBufNum;
441+ WaitFlag<HardEvent::FIX_M>(l0CFixToMFlagId + l0CBufId);
594 442 
595- inputLeftBuf.Wait<HardEvent::MTE2_MTE1>(); // 等待L1A
596- Buffer<BufferType::L0C> mm1ResL0C = mmL0CBuffers.Get();
597- mm1ResL0C.Wait<HardEvent::FIX_M>(); // 占用
598 MMParam param = {443 MMParam param = {
599- static_cast<uint32_t>(runInfo.mRealSize), // singleM444+ static_cast<uint32_t>(runInfo.mRealSize), // singleM
600- static_cast<uint32_t>(runInfo.s2RealSize), // singleN445+ static_cast<uint32_t>(runInfo.s2RealSize), // singleN
601- static_cast<uint32_t>(constInfo.dSize), // singleK446+ static_cast<uint32_t>(constInfo.dSize >> 1), // singleK
602- 0, // isLeftTranspose447+ 0, // isLeftTranspose
603- 1 // isRightTranspose448+ 1 // isRightTranspose
604 };449 };
605- MatmulK<Q_T, Q_T, T, s1BaseSize, s2BaseSize, dBaseMatmulSize, ABLayout::MK, ABLayout::KN>( // m,n不切，k切128450+ uint32_t curL1QBufId = l1QBufId;
606- inputLeftBuf.GetTensor<Q_T>(), inputRightBuf.GetTensor<Q_T>(), // mm1B直接用tensor的数据451+ if (unlikely(runInfo.s2LoopCount == 0)) {
607- mmL0ABuffers, mmL0BBuffers, mm1ResL0C.GetTensor<T>(), param);452+ WaitFlag<HardEvent::MTE2_MTE1>(l1QMte2ToMte1FlagId + curL1QBufId);
608- if (unlikely(runInfo.s2LoopCount == runInfo.s2LoopLimit)) {
609- inputLeftBuf.Set<HardEvent::MTE1_MTE2>(); // 释放L1A
610 }453 }
611 454 
612- mm1ResL0C.Set<HardEvent::M_FIX>(); // 通知455+ // m,n不切，k切128，mm1B直接用tensor的数据
613- mm1ResL0C.Wait<HardEvent::M_FIX>(); // 等待L0C456+ MatmulK<Q_T, Q_T, T, s1BaseSize, s2BaseSize, dBaseMatmulSize, ABLayout::MK, ABLayout::KN>(
457+ l1QTensor[curL1QBufId * BUFFER_SIZE_16K],
458+ curL1RightTensor,
459+ mmL0ABuffers, mmL0BBuffers,
460+ mmL0CTensor[BUFFER_SIZE_32K * l0CBufId],
461+ param);
462+ 
463+ curL1QBufId = (curL1QBufId + 1) % l1QBufNum;
464+ if (unlikely(runInfo.s2LoopCount == 0)) {
465+ WaitFlag<HardEvent::MTE2_MTE1>(l1QMte2ToMte1FlagId + curL1QBufId);
466+ }
467+ param.singleK = constInfo.dSize - param.singleK;
468+ param.isOutKFisrt = false;
469+ 
470+ // m,n不切，k切128, mm1B直接用tensor的数据
471+ MatmulK<Q_T, Q_T, T, s1BaseSize, s2BaseSize, dBaseMatmulSize, ABLayout::MK, ABLayout::KN>(
472+ l1QTensor[curL1QBufId * BUFFER_SIZE_16K],
473+ curL1RightTensor[(constInfo.dSize >> 1) * Align16Func(runInfo.s2RealSize)],
474+ mmL0ABuffers, mmL0BBuffers,
475+ mmL0CTensor[BUFFER_SIZE_32K * l0CBufId],
476+ param);
477+ 
478+ if (unlikely(runInfo.s2LoopCount == runInfo.s2LoopLimit)) {
479+ SetFlag<HardEvent::MTE1_MTE2>(l1QMte1ToMte2FlagId + l1QBufId);
480+ SetFlag<HardEvent::MTE1_MTE2>(l1QMte1ToMte2FlagId + curL1QBufId);
481+ l1QBufId = (l1QBufId + qHalfNum) % l1QBufNum;
482+ if (notLastTwoLoop) {
483+ CopyQGmToL1(runInfoNext, constInfo);
484+ }
485+ }
486+ 
487+ SetFlag<HardEvent::M_FIX>(l0CMToFixFlagId + l0CBufId);
488+ WaitFlag<HardEvent::M_FIX>(l0CMToFixFlagId + l0CBufId);
614 489 
615 outputBuf.WaitCrossCore();490 outputBuf.WaitCrossCore();
616- // L0C→UB491+ FixpipeParamsC310<CO2Layout::ROW_MAJOR> fixpipeParams; // L0C→UB
617- FixpipeParamsC310<CO2Layout::ROW_MAJOR> fixpipeParams;
618 // L0C上的bmm1结果矩阵N方向的size大小; 同mmadParams.n; 为什么要8个元素对齐(32B对齐) // 128492 // L0C上的bmm1结果矩阵N方向的size大小; 同mmadParams.n; 为什么要8个元素对齐(32B对齐) // 128
619 fixpipeParams.nSize = Align8Func(runInfo.s2RealSize);493 fixpipeParams.nSize = Align8Func(runInfo.s2RealSize);
620 // 有效数据不足16行，只需要输出部分行即可; L0C上的bmm1结果矩阵M方向的size大小(必须为偶数) // 128494 // 有效数据不足16行，只需要输出部分行即可; L0C上的bmm1结果矩阵M方向的size大小(必须为偶数) // 128
@@ -631,45 +505,45 @@ CSABlockCube<TEMPLATE_ARGS>::IterateBmm1CSA(Buffer<BufferType::UB, SyncType::CRO
631 fixpipeParams.params.dstNdStride = 0;505 fixpipeParams.params.dstNdStride = 0;
632 506 
633 // 将matmul结果从L0C搬运到UB507 // 将matmul结果从L0C搬运到UB
634- Fixpipe<T, T, PFA_CFG_ROW_MAJOR_UB>(outputBuf.template GetTensor<T>(), mm1ResL0C.GetTensor<T>(), fixpipeParams);508+ Fixpipe<T, T, PFA_CFG_ROW_MAJOR_UB>(outputBuf.template GetTensor<T>(),
635- mm1ResL0C.Set<HardEvent::FIX_M>(); // 释放L0C509+ mmL0CTensor[BUFFER_SIZE_32K * l0CBufId], fixpipeParams);
510+ SetFlag<HardEvent::FIX_M>(l0CFixToMFlagId + l0CBufId);
511+ l0CBufId ^= 1;
636 outputBuf.SetCrossCore();512 outputBuf.SetCrossCore();
637}513}
638 514 
639TEMPLATES_DEF_NO_DEFAULT515TEMPLATES_DEF_NO_DEFAULT
640-__aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::IterateBmm2CSA(516+__aicore__ inline void
641- Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,517+CSABlockCube<TEMPLATE_ARGS>::IterateBmm2Impl(Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &outputBuf,
642- BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputLeftBuffers,518+ BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputLeftBuffers,
643- Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &inputRightBuf, const RunInfo &runInfo,519+ const RunInfo &runInfo,
644- const ConstInfo &constInfo)520+ const ConstInfo &constInfo)
645{521{
646- Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> inputLeftBuf = inputLeftBuffers.Get(); // P直接用无需搬运522+ LocalTensor<Q_T> curL1RightTensor = l1RightTensor[runInfo.taskIdMod3 * rightBufSingleSize];
647- inputLeftBuf.WaitCrossCore();523+ Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> l1PBuffer = inputLeftBuffers.Get(); // P直接用无需搬运
524+ l1PBuffer.WaitCrossCore();
648 525 
649- Buffer<BufferType::L0C> mm2ResL0C = mmL0CBuffers.Get();526+ WaitFlag<HardEvent::FIX_M>(l0CFixToMFlagId + l0CBufId);
650- mm2ResL0C.Wait<HardEvent::FIX_M>(); // 占用
651 MMParam param = {527 MMParam param = {
652- static_cast<uint32_t>(runInfo.mRealSize), // singleM 64528+ static_cast<uint32_t>(runInfo.mRealSize), // singleM
653 static_cast<uint32_t>(constInfo.dSizeV), // singleN 512529 static_cast<uint32_t>(constInfo.dSizeV), // singleN 512
654 static_cast<uint32_t>(runInfo.s2RealSize), // singleK 128530 static_cast<uint32_t>(runInfo.s2RealSize), // singleK 128
655 0, // isLeftTranspose531 0, // isLeftTranspose
656 0 // isRightTranspose532 0 // isRightTranspose
657 };533 };
658 MatmulN<Q_T, Q_T, T, s1BaseSize, s2BaseSize, dBaseMatmulSize, ABLayout::MK, ABLayout::KN>(534 MatmulN<Q_T, Q_T, T, s1BaseSize, s2BaseSize, dBaseMatmulSize, ABLayout::MK, ABLayout::KN>(
659- inputLeftBuf.GetTensor<Q_T>(), inputRightBuf.GetTensor<Q_T>(), mmL0ABuffers, mmL0BBuffers,535+ l1PBuffer.template GetTensor<Q_T>(),
660- mm2ResL0C.GetTensor<T>(), param);536+ curL1RightTensor,
537+ mmL0ABuffers, mmL0BBuffers,
538+ mmL0CTensor[BUFFER_SIZE_32K * l0CBufId],
539+ param);
661 540 
662- if constexpr (TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE ||541+ SetFlag<HardEvent::M_FIX>(l0CMToFixFlagId + l0CBufId);
663- TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||542+ WaitFlag<HardEvent::M_FIX>(l0CMToFixFlagId + l0CBufId);
664- TEMPLATE_MODE == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) {543+ SetFlag<HardEvent::MTE1_MTE2>(l1KMte1ToMte2FlagId + l1KMatmul2BufId);
665- SetFlag<HardEvent::MTE1_MTE2>(mte2ToMte1Id[runInfo.taskIdMod3]);544+ l1KMatmul2BufId = (l1KMatmul2BufId + 1) % l1KBufNum;
666- }
667 545 
668- mm2ResL0C.Set<HardEvent::M_FIX>(); // 通知546+ outputBuf.WaitCrossCore(); // 占用
669- mm2ResL0C.Wait<HardEvent::M_FIX>(); // 等待
670- 
671- // 占用
672- outputBuf.WaitCrossCore();
673 // L0C→UB;FixpipeParamsM300:L0C→UB547 // L0C→UB;FixpipeParamsM300:L0C→UB
674 FixpipeParamsC310<CO2Layout::ROW_MAJOR> fixpipeParams;548 FixpipeParamsC310<CO2Layout::ROW_MAJOR> fixpipeParams;
675 // L0C上的bmm1结果矩阵N方向的size大小, 分档计算且vector2中通过mask筛选出实际有效值549 // L0C上的bmm1结果矩阵N方向的size大小, 分档计算且vector2中通过mask筛选出实际有效值
@@ -683,9 +557,10 @@ __aicore__ inline void CSABlockCube<TEMPLATE_ARGS>::IterateBmm2CSA(
683 fixpipeParams.params.ndNum = 1;557 fixpipeParams.params.ndNum = 1;
684 fixpipeParams.params.srcNdStride = 0;558 fixpipeParams.params.srcNdStride = 0;
685 fixpipeParams.params.dstNdStride = 0;559 fixpipeParams.params.dstNdStride = 0;
686- Fixpipe<T, T, PFA_CFG_ROW_MAJOR_UB>(outputBuf.template GetTensor<T>(), mm2ResL0C.GetTensor<T>(),560+ Fixpipe<T, T, PFA_CFG_ROW_MAJOR_UB>(outputBuf.template GetTensor<T>(),
687- fixpipeParams); // 将matmul结果从L0C搬运到UB561+ mmL0CTensor[BUFFER_SIZE_32K * l0CBufId], fixpipeParams); // 将matmul结果从L0C搬运到UB
688- mm2ResL0C.Set<HardEvent::FIX_M>(); // 释放562+ SetFlag<HardEvent::FIX_M>(l0CFixToMFlagId + l0CBufId);
563+ l0CBufId ^= 1;
689 564 
690 outputBuf.SetCrossCore();565 outputBuf.SetCrossCore();
691}566}
@@ -694,10 +569,7 @@ TEMPLATES_DEF
694class CSABlockCubeDummy {569class CSABlockCubeDummy {
695public:570public:
696 __aicore__ inline CSABlockCubeDummy(){};571 __aicore__ inline CSABlockCubeDummy(){};
697- __aicore__ inline void InitCubeBlock(TPipe *pipe, BufferManager<BufferType::L1> &l1BufferManager,572+ __aicore__ inline void InitCubeBlock(TPipe *pipe, __gm__ uint8_t *query) {}
698- __gm__ uint8_t *query)
699- {
700- }
701 __aicore__ inline void InitCubeInput(__gm__ uint8_t *oriKv, __gm__ uint8_t *cmpKv, __gm__ uint8_t *cmpSparseIndices,573 __aicore__ inline void InitCubeInput(__gm__ uint8_t *oriKv, __gm__ uint8_t *cmpKv, __gm__ uint8_t *cmpSparseIndices,
702 __gm__ uint8_t *oriBlockTable, __gm__ uint8_t *cmpBlockTable,574 __gm__ uint8_t *oriBlockTable, __gm__ uint8_t *cmpBlockTable,
703 __gm__ uint8_t *sequsedQ, __gm__ uint8_t *cuSeqlensQ,575 __gm__ uint8_t *sequsedQ, __gm__ uint8_t *cuSeqlensQ,
@@ -715,11 +587,11 @@ struct CubeBlockTraits; // 声明
715#define GEN_TRAIT_TYPE(name, ...) using name##_TRAITS = name;587#define GEN_TRAIT_TYPE(name, ...) using name##_TRAITS = name;
716#define GEN_TRAIT_CONST(name, type, ...) static constexpr type name##Traits = name;588#define GEN_TRAIT_CONST(name, type, ...) static constexpr type name##Traits = name;
717 589 
718-#define DEFINE_CUBE_BLOCK_TRAITS(CUBE_BLOCK_CLASS) \590+#define DEFINE_CUBE_BLOCK_TRAITS(CUBE_BLOCK_CLASS) \
719- TEMPLATES_DEF_NO_DEFAULT \591+ TEMPLATES_DEF_NO_DEFAULT \
720- struct CubeBlockTraits<CUBE_BLOCK_CLASS<TEMPLATE_ARGS>> { \592+ struct CubeBlockTraits<CUBE_BLOCK_CLASS<TEMPLATE_ARGS>> { \
721- CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_TRAIT_TYPE) \593+ CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_TRAIT_TYPE) \
722- CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_TRAIT_CONST) \594+ CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_TRAIT_CONST) \
723 }595 }
724 596 
725DEFINE_CUBE_BLOCK_TRAITS(CSABlockCube);597DEFINE_CUBE_BLOCK_TRAITS(CSABlockCube);
@@ -728,8 +600,8 @@ DEFINE_CUBE_BLOCK_TRAITS(CSABlockCubeDummy);
728// /* 生成Arg Traits, kernel中只需要调用ARGS_TRAITS就可以获取所有CubeBlock中的模板参数 */600// /* 生成Arg Traits, kernel中只需要调用ARGS_TRAITS就可以获取所有CubeBlock中的模板参数 */
729#define GEN_ARGS_TYPE(name, ...) using name = typename CubeBlockTraits<CubeBlockType>::name##_TRAITS;601#define GEN_ARGS_TYPE(name, ...) using name = typename CubeBlockTraits<CubeBlockType>::name##_TRAITS;
730#define GEN_ARGS_CONST(name, type, ...) static constexpr type name = CubeBlockTraits<CubeBlockType>::name##Traits;602#define GEN_ARGS_CONST(name, type, ...) static constexpr type name = CubeBlockTraits<CubeBlockType>::name##Traits;
731-#define ARGS_TRAITS \603+#define ARGS_TRAITS \
732- CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_ARGS_TYPE) \604+ CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_ARGS_TYPE) \
733 CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_ARGS_CONST)605 CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_ARGS_CONST)
734} // namespace SMLAKernel606} // namespace SMLAKernel
735#endif // FLASH_ATTENTION_SCORE_BLOCK_CUBE_H_607#endif // FLASH_ATTENTION_SCORE_BLOCK_CUBE_H_
Mattention/sparse_flash_mla/op_kernel/arch35/sparse_flash_mla_csa_block_vector.h+1013-182
@@ -80,6 +80,14 @@ public:
80 static constexpr uint32_t dTemplateAlign64 = Align64Func(dVTemplateType);80 static constexpr uint32_t dTemplateAlign64 = Align64Func(dVTemplateType);
81 static constexpr float R0 = 1.0f;81 static constexpr float R0 = 1.0f;
82 82 
83+ // GetKVPhyAddrForKvType 事件 ID (手动管理, 不通过 AllocEventID)
84+ static constexpr uint32_t blkTableVToMte2FlagId = 3; // block table V→MTE2
85+ static constexpr uint32_t sparseIdxVToMte2FlagId = 4; // sparse indices V→MTE2
86+ static constexpr uint32_t phyAddrVToMte3FlagId = 5; // phy addr V→MTE3
87+ static constexpr uint32_t sparseIdxMte2ToVFlagId = 6; // sparse indices MTE2→V
88+ static constexpr uint32_t phyAddrMte3ToVFlagId = 7; // phy addr MTE3→V
89+ static constexpr uint32_t blkTableMte2ToVFlagId = 8; // block table MTE2→V
90+ 
83 // ==================== Functions ======================91 // ==================== Functions ======================
84 __aicore__ inline CSABlockVec(){};92 __aicore__ inline CSABlockVec(){};
85 __aicore__ inline void InitVecBlock(TPipe *pipe, __gm__ uint8_t *cuSeqlensQ, __gm__ uint8_t *cuSeqlensOriKv,93 __aicore__ inline void InitVecBlock(TPipe *pipe, __gm__ uint8_t *cuSeqlensQ, __gm__ uint8_t *cuSeqlensOriKv,
@@ -121,13 +129,12 @@ public:
121 __gm__ uint8_t *sequsedOriKv, __gm__ uint8_t *sequsedCmpKv,129 __gm__ uint8_t *sequsedOriKv, __gm__ uint8_t *sequsedCmpKv,
122 __gm__ uint8_t *cmpResidualKv);130 __gm__ uint8_t *cmpResidualKv);
123 __aicore__ inline void InitOutputSingleCore(ConstInfo &constInfo);131 __aicore__ inline void InitOutputSingleCore(ConstInfo &constInfo);
124- __aicore__ inline void ProcessVec0(Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &outputL1,132+ __aicore__ inline void ProcessVec0(Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,
125- Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,133+ const RunInfo &runInfo, ConstInfo &constInfo);
126- const RunInfo &runInfo, ConstInfo &constInfo, int32_t startPos);
127 __aicore__ inline void ProcessVec1(Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &outputBuf,134 __aicore__ inline void ProcessVec1(Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &outputBuf,
128- Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &bmm1ResBuf,135+ Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH> &bmm1ResBuf,
129- RunInfo &runInfo,136+ RunInfo &runInfo,
130- ConstInfo &constInfo);137+ ConstInfo &constInfo);
131 __aicore__ inline void InitS2SplitStaging(138 __aicore__ inline void InitS2SplitStaging(
132 Buffer<BufferType::GM, SyncType::INNER_CORE_SYNC> &fdStaging)139 Buffer<BufferType::GM, SyncType::INNER_CORE_SYNC> &fdStaging)
133 {140 {
@@ -149,12 +156,23 @@ public:
149 __aicore__ inline void ProcessFlashDecode(FdRunInfo &fdRunInfo, ConstInfo &constInfo);156 __aicore__ inline void ProcessFlashDecode(FdRunInfo &fdRunInfo, ConstInfo &constInfo);
150 using mm2ResPos = Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH>;157 using mm2ResPos = Buffer<BufferType::UB, SyncType::CROSS_CORE_SYNC_BOTH>;
151 __aicore__ inline void ProcessVec2(mm2ResPos &bmm2ResBuf, RunInfo &runInfo, ConstInfo &constInfo);158 __aicore__ inline void ProcessVec2(mm2ResPos &bmm2ResBuf, RunInfo &runInfo, ConstInfo &constInfo);
159+ __aicore__ inline void GetKVPhyAddr(uint32_t hasLoad, uint32_t bN2StartIdx, uint32_t bN2EndIdx,
160+ uint32_t gS1StartIdx, uint32_t nextGs1Idx, bool hasActualSeqQlen,
161+ bool hasCuSeqlensQ, bool hasActualSeqOriKvlen, bool hasCuSeqlensOriKv,
162+ GlobalTensor<int32_t> actualSeqOriKvlenGm,
163+ GlobalTensor<int32_t> cuSeqlensOriKvGm, GlobalTensor<int32_t> oriTopkLengthGm,
164+ bool hasActualSeqCmpKvlen, bool hasCuSeqlensCmpKv,
165+ GlobalTensor<int32_t> actualSeqCmpKvlenGm,
166+ GlobalTensor<int32_t> cuSeqlensCmpKvGm, GlobalTensor<int32_t> cmpTopkLengthGm,
167+ GlobalTensor<int32_t> cmpResidualKvGm, GlobalTensor<int32_t> actualSeqQlenGm,
168+ GlobalTensor<int32_t> cuSeqlensQGm, __gm__ uint8_t *workspace,
169+ ConstInfo &constInfo);
152 170 
153private:171private:
154 template <bool UPDATE>172 template <bool UPDATE>
155 __aicore__ inline void ComputeVec1Softmax(LocalTensor<Q_T> &stage1CastTensor, LocalTensor<T> &mmRes,173 __aicore__ inline void ComputeVec1Softmax(LocalTensor<Q_T> &stage1CastTensor, LocalTensor<T> &mmRes,
156- LocalTensor<float> &sumUb, LocalTensor<float> &maxUb, LocalTensor<T> &apiTmpBuffer,174+ LocalTensor<float> &sumUb, LocalTensor<float> &maxUb, LocalTensor<T> &apiTmpBuffer,
157- RunInfo &runInfo, ConstInfo &constInfo);175+ RunInfo &runInfo, ConstInfo &constInfo);
158 __aicore__ inline void InitVec1SoftmaxFromSinks(176 __aicore__ inline void InitVec1SoftmaxFromSinks(
159 LocalTensor<float> &sumUb, LocalTensor<float> &maxUb, RunInfo &runInfo, ConstInfo &constInfo);177 LocalTensor<float> &sumUb, LocalTensor<float> &maxUb, RunInfo &runInfo, ConstInfo &constInfo);
160 __aicore__ inline void CopyVec1ResultToL1(178 __aicore__ inline void CopyVec1ResultToL1(
@@ -216,25 +234,49 @@ private:
216 return stagingMOffset;234 return stagingMOffset;
217 }235 }
218 236 
219- __aicore__ inline void ProcessSparseKv(Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &outputL1,237+ __aicore__ inline void ProcessSparseKv(Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,
220- Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,238+ const RunInfo &runInfo, ConstInfo &constInfo);
221- const RunInfo &runInfo, ConstInfo &constInfo, int32_t startPos);
222 __aicore__ inline void CalSparseCalSize(const RunInfo &runInfo, ConstInfo &constInfo);239 __aicore__ inline void CalSparseCalSize(const RunInfo &runInfo, ConstInfo &constInfo);
223 __aicore__ inline int64_t GetkeyOffset(int64_t s2Idx, const RunInfo &runInfo, ConstInfo &constInfo);240 __aicore__ inline int64_t GetkeyOffset(int64_t s2Idx, const RunInfo &runInfo, ConstInfo &constInfo);
224- __aicore__ inline void GetRealCmpS2Idx(int64_t &token0Idx, int64_t &token1Idx, int64_t s2IdxInBase,241+ __aicore__ inline void GetRealCmpS2Idx(int64_t *tokenData, int64_t s2IdxInBase,
225 const RunInfo &runInfo, ConstInfo &constInfo);242 const RunInfo &runInfo, ConstInfo &constInfo);
226- __aicore__ inline uint32_t CopyInKvSparse(LocalTensor<KV_T> kvInUb, int64_t startRow, int64_t token0Idx,243+ __aicore__ inline uint32_t CopyInKvSparse(LocalTensor<KV_T> kvInUb, int64_t startRow, int64_t *tokenData,
227- int64_t token1Idx, const RunInfo &runInfo, ConstInfo &constInfo);244+ const RunInfo &runInfo, ConstInfo &constInfo);
228 __aicore__ inline void CopyToOutUb(LocalTensor<Q_T> kvNzUb, LocalTensor<KV_T> srcTensor, int64_t dealRow,245 __aicore__ inline void CopyToOutUb(LocalTensor<Q_T> kvNzUb, LocalTensor<KV_T> srcTensor, int64_t dealRow,
229 ConstInfo &constInfo);246 ConstInfo &constInfo);
230- __aicore__ inline void CopyOutKvUb2L1(Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &outputL1,
231- LocalTensor<Q_T> kvNzOutUb, int64_t dealRow, int64_t s2StartIdx,
232- const RunInfo &runInfo, ConstInfo &constInfo);
233 __aicore__ inline void CopyOutKvUb2Gm(Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,247 __aicore__ inline void CopyOutKvUb2Gm(Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,
234 LocalTensor<Q_T> kvOutUb, int64_t dealRow, int64_t s2StartIdx,248 LocalTensor<Q_T> kvOutUb, int64_t dealRow, int64_t s2StartIdx,
235 const RunInfo &runInfo, ConstInfo &constInfo);249 const RunInfo &runInfo, ConstInfo &constInfo);
236 __aicore__ inline void CopyInSingleKv(LocalTensor<KV_T> kvInUb, int64_t startRow, int64_t keyOffset,250 __aicore__ inline void CopyInSingleKv(LocalTensor<KV_T> kvInUb, int64_t startRow, int64_t keyOffset,
237 ConstInfo &constInfo);251 ConstInfo &constInfo);
252+ __aicore__ inline void GetRealS2Addr(int64_t *tokenData, int64_t s2IdxInBase, const RunInfo &runInfo,
253+ ConstInfo &constInfo);
254+ __aicore__ inline void GetKVPhyAddrForKvType(
255+ uint32_t bN2StartIdx, uint32_t bN2EndIdx, uint32_t gS1StartIdx, uint32_t nextGs1Idx, bool hasActualSeqQlen,
256+ bool hasCuSeqlensQ, bool hasActualSeqKvlen, bool hasCuSeqlensKv, GlobalTensor<int32_t> actualSeqQlenGm,
257+ GlobalTensor<int32_t> cuSeqlensQGm, GlobalTensor<int32_t> actualSeqKvlenGm, GlobalTensor<int32_t> cuSeqlensKvGm,
258+ GlobalTensor<int32_t> topkLengthGm, GlobalTensor<int32_t> cmpResidualKvGm, ConstInfo &constInfo,
259+ GlobalTensor<int32_t> &blockTableGm, GlobalTensor<int32_t> &sparseIndicesGm, GlobalTensor<uint32_t> &phyAddrGm,
260+ uint32_t kvStride, uint32_t blockSize, uint32_t maxBlockNumPerBatch, uint32_t sparseBlockCount,
261+ uint32_t alignedSparseBlockCount, bool isOriKv);
262+ __aicore__ inline int32_t GetSeqLen(int32_t bIdx, bool hasActualSeq, bool hasCuSeqlens,
263+ GlobalTensor<int32_t> &actualSeqGm, GlobalTensor<int32_t> &cuSeqlensGm,
264+ int64_t defaultSize);
265+ __aicore__ inline int32_t CalcCurValidS2(uint32_t bIdx, int32_t s1Idx, int32_t actualS1Size, bool isOriKv,
266+ bool hasActualSeqKvlen, bool hasCuSeqlensKv,
267+ GlobalTensor<int32_t> &cuSeqlensQGm,
268+ GlobalTensor<int32_t> &actualSeqKvlenGm,
269+ GlobalTensor<int32_t> &cuSeqlensKvGm, GlobalTensor<int32_t> &topkLengthGm,
270+ GlobalTensor<int32_t> &cmpResidualKvGm, ConstInfo &constInfo,
271+ int32_t sparseBlockCount);
272+ __aicore__ inline void CopyPhyAddrToGm(LocalTensor<uint32_t> kvPhyAddrUb, int64_t bS1Idx, int64_t s1Idx,
273+ int64_t validS2, int64_t alignNum, GlobalTensor<uint32_t> &phyAddrGm,
274+ uint32_t alignedSparseBlockCount);
275+ __aicore__ inline void CopyPaTableToUb(LocalTensor<int32_t> blkTableUb, int64_t bIdx,
276+ GlobalTensor<int32_t> &blockTableGm, uint32_t maxBlockNumPerBatch);
277+ __aicore__ inline void CopySparseIdxToUb(LocalTensor<int32_t> sparseIdxUb, int64_t bS1Idx, int64_t s1Idx,
278+ int64_t validS2, GlobalTensor<int32_t> &sparseIndicesGm,
279+ uint32_t sparseBlockCount);
238 /* VEC2_RES_T 表示bmm2ResUb当前的类型，VEC2_RES_T = Q_T那么不需要做Cast。另外，无效行场景当前默认需要做Cast */280 /* VEC2_RES_T 表示bmm2ResUb当前的类型，VEC2_RES_T = Q_T那么不需要做Cast。另外，无效行场景当前默认需要做Cast */
239 template <typename VEC2_RES_T>281 template <typename VEC2_RES_T>
240 __aicore__ inline void Bmm2DataCopyOut(RunInfo &runInfo, ConstInfo &constInfo, LocalTensor<VEC2_RES_T> &vec2ResUb,282 __aicore__ inline void Bmm2DataCopyOut(RunInfo &runInfo, ConstInfo &constInfo, LocalTensor<VEC2_RES_T> &vec2ResUb,
@@ -247,7 +289,7 @@ private:
247 __aicore__ inline void GetExtremeValue(T &negativeScalar);289 __aicore__ inline void GetExtremeValue(T &negativeScalar);
248 __aicore__ inline void InitSinksBuffer(ConstInfo &constInfo);290 __aicore__ inline void InitSinksBuffer(ConstInfo &constInfo);
249 __aicore__ inline void ReduceIntraBlockAndStage(RunInfo &runInfo, ConstInfo &constInfo,291 __aicore__ inline void ReduceIntraBlockAndStage(RunInfo &runInfo, ConstInfo &constInfo,
250- LocalTensor<T> &vec2ResUb, LocalTensor<T> &partialTmpUb);292+ LocalTensor<T> &vec2ResUb, LocalTensor<T> &partialTmpUb);
251 293 
252 TPipe *tPipe;294 TPipe *tPipe;
253 295 
@@ -270,6 +312,8 @@ private:
270 GlobalTensor<int32_t> actualSeqLengthsKVGm;312 GlobalTensor<int32_t> actualSeqLengthsKVGm;
271 GlobalTensor<int32_t> actualSeqLengthsCmpKVGm;313 GlobalTensor<int32_t> actualSeqLengthsCmpKVGm;
272 GlobalTensor<int32_t> cmpResidualKVGm;314 GlobalTensor<int32_t> cmpResidualKVGm;
315+ GlobalTensor<uint32_t> oriKvPhyAddrGm;
316+ GlobalTensor<uint32_t> cmpKvPhyAddrGm;
273 317 
274 TBuf<> commonTBuf; // common的复用空间318 TBuf<> commonTBuf; // common的复用空间
275 TBuf<> sinksBuf;319 TBuf<> sinksBuf;
@@ -327,21 +371,25 @@ private:
327};371};
328 372 
329TEMPLATES_DEF_NO_DEFAULT373TEMPLATES_DEF_NO_DEFAULT
330-__aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::GetRealCmpS2Idx(int64_t &token0Idx, int64_t &token1Idx,374+__aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::GetRealCmpS2Idx(int64_t *tokenData, int64_t s2IdxInBase,
331- int64_t s2IdxInBase, const RunInfo &runInfo,375+ const RunInfo &runInfo, ConstInfo &constInfo)
332- ConstInfo &constInfo)
333{376{
334 int64_t sparseBlockCount = 0;377 int64_t sparseBlockCount = 0;
335 int64_t curS2LoopCnt = runInfo.s2LoopCount;378 int64_t curS2LoopCnt = runInfo.s2LoopCount;
336- // CSA、ORI_SPARSE、ORI_CMP_SPARSE均可通过runInfo.isCmp判断379+ if constexpr (TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE) {
337- if (runInfo.isCmp) {
338 sparseBlockCount = constInfo.cmpSparseBlockCount;380 sparseBlockCount = constInfo.cmpSparseBlockCount;
339 curS2LoopCnt -= runInfo.oriKvLoopEndIdx;381 curS2LoopCnt -= runInfo.oriKvLoopEndIdx;
340- } else {382+ } else if constexpr (TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE) {
341 sparseBlockCount = constInfo.oriSparseBlockCount;383 sparseBlockCount = constInfo.oriSparseBlockCount;
384+ } else if constexpr (TEMPLATE_MODE == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) {
385+ if (runInfo.isCmp) {
386+ sparseBlockCount = constInfo.cmpSparseBlockCount;
387+ curS2LoopCnt -= runInfo.oriKvLoopEndIdx;
388+ } else {
389+ sparseBlockCount = constInfo.oriSparseBlockCount;
390+ }
342 }391 }
343- 392+ uint64_t topkBS1Idx = 0;
344- int64_t topkBS1Idx = 0;
345 if constexpr (LAYOUT_T == SMLA_LAYOUT::TND) {393 if constexpr (LAYOUT_T == SMLA_LAYOUT::TND) {
346 uint64_t actualSeqQPrefixSum = cuSeqlensQGm.GetValue(runInfo.boIdx);394 uint64_t actualSeqQPrefixSum = cuSeqlensQGm.GetValue(runInfo.boIdx);
347 topkBS1Idx += (actualSeqQPrefixSum + runInfo.s1oIdx) * sparseBlockCount; // T, N2(1), K395 topkBS1Idx += (actualSeqQPrefixSum + runInfo.s1oIdx) * sparseBlockCount; // T, N2(1), K
@@ -350,17 +398,49 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::GetRealCmpS2Idx(int64_t &toke
350 runInfo.boIdx * constInfo.s1Size * sparseBlockCount + runInfo.s1oIdx * sparseBlockCount; // B, S1, N2(1), K398 runInfo.boIdx * constInfo.s1Size * sparseBlockCount + runInfo.s1oIdx * sparseBlockCount; // B, S1, N2(1), K
351 }399 }
352 400 
353- int64_t topkKIdx = s2IdxInBase + curS2LoopCnt * constInfo.s2BaseSize;401+ uint64_t topkKIdx = s2IdxInBase + curS2LoopCnt * constInfo.s2BaseSize;
354- if (unlikely(topkKIdx >= sparseBlockCount)) {402+ for (uint64_t i = 0; i < 8; ++i) {
355- token0Idx = -1;403+ uint64_t idx = topkBS1Idx + runInfo.s2StartIdx + topkKIdx + i;
356- } else {404+ if (likely((topkKIdx + i < sparseBlockCount) && (s2IdxInBase + i < sparseS2End))) {
357- token0Idx = sparseIndicesGm.GetValue(topkBS1Idx + topkKIdx + runInfo.s2StartIdx);405+ tokenData[i] = sparseIndicesGm.GetValue(idx);
406+ } else {
407+ break;
408+ }
358 }409 }
359- topkKIdx += 1;410+}
360- if (unlikely((topkKIdx >= sparseBlockCount) || (s2IdxInBase + 1 >= sparseS2End))) {411+ 
361- token1Idx = -1;412+TEMPLATES_DEF_NO_DEFAULT
413+__aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::GetRealS2Addr(int64_t *tokenData, int64_t s2IdxInBase,
414+ const RunInfo &runInfo, ConstInfo &constInfo)
415+{
416+ uint32_t sparseBlockCount = runInfo.isCmp ? constInfo.cmpSparseBlockCount : constInfo.oriSparseBlockCount;
417+ uint32_t alignedSparseBlockCount =
418+ runInfo.isCmp ? constInfo.alignedCmpSparseBlockCount : constInfo.alignedOriSparseBlockCount;
419+ int64_t curS2LoopCnt = runInfo.s2LoopCount;
420+ GlobalTensor<int64_t> phyAddrGm64;
421+ if (runInfo.isCmp) {
422+ curS2LoopCnt -= runInfo.oriKvLoopEndIdx;
423+ phyAddrGm64 = cmpKvPhyAddrGm.template ReinterpretCast<int64_t>();
362 } else {424 } else {
363- token1Idx = sparseIndicesGm.GetValue(topkBS1Idx + topkKIdx + runInfo.s2StartIdx);425+ phyAddrGm64 = oriKvPhyAddrGm.template ReinterpretCast<int64_t>();
426+ }
427+ 
428+ uint64_t topkBS1Idx = 0;
429+ if constexpr (LAYOUT_T == SMLA_LAYOUT::TND) {
430+ uint64_t actualSeqQPrefixSum = cuSeqlensQGm.GetValue(runInfo.boIdx);
431+ topkBS1Idx += (actualSeqQPrefixSum + runInfo.s1oIdx) * alignedSparseBlockCount;
432+ } else {
433+ topkBS1Idx +=
434+ runInfo.boIdx * constInfo.s1Size * alignedSparseBlockCount + runInfo.s1oIdx * alignedSparseBlockCount;
435+ }
436+ uint64_t topkKIdx = s2IdxInBase + curS2LoopCnt * constInfo.s2BaseSize;
437+ for (uint64_t i = 0; i < 8; ++i) {
438+ uint64_t idx = topkBS1Idx + runInfo.s2StartIdx + topkKIdx + i;
439+ if (likely((topkKIdx + i < sparseBlockCount) && (s2IdxInBase + i < sparseS2End))) {
440+ tokenData[i] = phyAddrGm64.GetValue(idx);
441+ } else {
442+ break;
443+ }
364 }444 }
365}445}
366 446 
@@ -418,42 +498,53 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::CopyInSingleKv(LocalTensor<KV
418 498 
419TEMPLATES_DEF_NO_DEFAULT499TEMPLATES_DEF_NO_DEFAULT
420__aicore__ inline uint32_t CSABlockVec<TEMPLATE_ARGS>::CopyInKvSparse(LocalTensor<KV_T> kvInUb, int64_t startRow,500__aicore__ inline uint32_t CSABlockVec<TEMPLATE_ARGS>::CopyInKvSparse(LocalTensor<KV_T> kvInUb, int64_t startRow,
421- int64_t token0Idx, int64_t token1Idx,501+ int64_t *tokenData, const RunInfo &runInfo,
422- const RunInfo &runInfo, ConstInfo &constInfo)502+ ConstInfo &constInfo)
423{503{
424- int64_t keyOffset0 = GetkeyOffset(token0Idx, runInfo, constInfo);504+ uint32_t dealRow = 0;
425- int64_t keyOffset1 = GetkeyOffset(token1Idx, runInfo, constInfo);505+ for (uint32_t i = 0; i < 8; i += 2) {
426- if (unlikely(keyOffset0 < 0 && keyOffset1 < 0)) {506+ int64_t keyOffset0;
427- return 0;507+ int64_t keyOffset1;
428- }508+ if constexpr (IS_VEC_S2PHYADDR) {
429- int64_t combineBytes = constInfo.dSizeVInput * sizeof(KV_T);509+ keyOffset0 = tokenData[i];
430- int64_t keySrcStride =510+ keyOffset1 = tokenData[i + 1];
431- (keyOffset0 > keyOffset1 ? (keyOffset0 - keyOffset1) : (keyOffset1 - keyOffset0)) * sizeof(KV_T) - combineBytes;511+ } else {
432- if (unlikely(keyOffset1 < 0)) {512+ keyOffset0 = GetkeyOffset(tokenData[i], runInfo, constInfo);
433- CopyInSingleKv(kvInUb, startRow, keyOffset0, constInfo);513+ keyOffset1 = GetkeyOffset(tokenData[i + 1], runInfo, constInfo);
434- } else if (unlikely(keySrcStride >= INT32_MAX || keySrcStride < 0) || constInfo.sparseBlockSize > 1) {
435- // stride溢出、stride为负数、s2超长等异常场景，还原成2条搬运指令
436- CopyInSingleKv(kvInUb, startRow, keyOffset0, constInfo);
437- CopyInSingleKv(kvInUb, startRow + 1, keyOffset1, constInfo);
438- } else {
439- DataCopyExtParams intriParams;
440- intriParams.blockCount = (keyOffset0 >= 0) + (keyOffset1 >= 0);
441- intriParams.blockLen = combineBytes;
442- intriParams.dstStride = 0;
443- intriParams.srcStride = keySrcStride;
444- DataCopyPadExtParams<KV_T> padParams;
445- padParams.isPad = true;
446- padParams.leftPadding = 0;
447- padParams.rightPadding = (CeilAlign(combineBytes, BUFFER_SIZE_BYTE_32B) - combineBytes) / sizeof(KV_T);
448- padParams.paddingValue = 0;
449- 
450- int64_t keyOffset = keyOffset0 > -1 ? keyOffset0 : keyOffset1;
451- if (keyOffset1 > -1 && keyOffset1 < keyOffset0) {
452- keyOffset = keyOffset1;
453 }514 }
454- DataCopyPad(kvInUb[startRow * constInfo.dSize], keyGm[keyOffset], intriParams, padParams);515+ if (unlikely(keyOffset0 < 0 && keyOffset1 < 0)) {
516+ return dealRow;
517+ }
518+ int64_t combineBytes = constInfo.dSizeVInput * sizeof(KV_T);
519+ int64_t keySrcStride =
520+ (keyOffset0 > keyOffset1 ? (keyOffset0 - keyOffset1) : (keyOffset1 - keyOffset0)) * sizeof(KV_T) -
521+ combineBytes;
522+ if (unlikely(keySrcStride >= INT32_MAX || keySrcStride < 0) || constInfo.sparseBlockSize > 1) {
523+ // stride溢出、stride为负数、s2超长等异常场景，还原成2条搬运指令
524+ CopyInSingleKv(kvInUb, startRow, keyOffset0, constInfo);
525+ CopyInSingleKv(kvInUb, startRow + 1, keyOffset1, constInfo);
526+ } else {
527+ DataCopyExtParams intriParams;
528+ intriParams.blockCount = (keyOffset0 >= 0) + (keyOffset1 >= 0);
529+ intriParams.blockLen = combineBytes;
530+ intriParams.dstStride = 0;
531+ intriParams.srcStride = keySrcStride;
532+ DataCopyPadExtParams<KV_T> padParams;
533+ padParams.isPad = true;
534+ padParams.leftPadding = 0;
535+ padParams.rightPadding = (CeilAlign(combineBytes, BUFFER_SIZE_BYTE_32B) - combineBytes) / sizeof(KV_T);
536+ padParams.paddingValue = 0;
537+ 
538+ int64_t keyOffset = keyOffset0 > -1 ? keyOffset0 : keyOffset1;
539+ if (keyOffset1 > -1 && keyOffset1 < keyOffset0) {
540+ keyOffset = keyOffset1;
541+ }
542+ DataCopyPad(kvInUb[startRow * constInfo.dSize], keyGm[keyOffset], intriParams, padParams);
543+ }
544+ dealRow += (keyOffset0 >= 0) + (keyOffset1 >= 0);
545+ startRow += 2;
455 }546 }
456- return (keyOffset0 > -1) + (keyOffset1 > -1);547+ return dealRow;
457}548}
458 549 
459TEMPLATES_DEF_NO_DEFAULT550TEMPLATES_DEF_NO_DEFAULT
@@ -464,23 +555,6 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::CopyToOutUb(LocalTensor<Q_T>
464 DataCopy(kvOutUb, kvNdUb, dealRow * constInfo.dSize);555 DataCopy(kvOutUb, kvNdUb, dealRow * constInfo.dSize);
465}556}
466 557 
467-TEMPLATES_DEF_NO_DEFAULT
468-__aicore__ inline void
469-CSABlockVec<TEMPLATE_ARGS>::CopyOutKvUb2L1(Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &outputL1,
470- LocalTensor<Q_T> kvNzOutUb, int64_t dealRow, int64_t s2StartIdx,
471- const RunInfo &runInfo, ConstInfo &constInfo)
472-{
473- uint64_t blockElementNum = 16;
474- DataCopyParams dataCopyParams;
475- dataCopyParams.blockCount = constInfo.dSize / blockElementNum;
476- dataCopyParams.blockLen = dealRow;
477- dataCopyParams.srcGap = blockElementNum + 1 - dealRow;
478- dataCopyParams.dstGap = Align16Func(runInfo.s2RealSize) - dealRow;
479- 
480- LocalTensor<Q_T> dst = outputL1.GetTensor<Q_T>();
481- DataCopy(dst[s2StartIdx * 16], kvNzOutUb, dataCopyParams);
482-}
483- 
484TEMPLATES_DEF_NO_DEFAULT558TEMPLATES_DEF_NO_DEFAULT
485__aicore__ inline void559__aicore__ inline void
486CSABlockVec<TEMPLATE_ARGS>::CopyOutKvUb2Gm(Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,560CSABlockVec<TEMPLATE_ARGS>::CopyOutKvUb2Gm(Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,
@@ -527,16 +601,11 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::CalSparseCalSize(const RunInf
527 601 
528TEMPLATES_DEF_NO_DEFAULT602TEMPLATES_DEF_NO_DEFAULT
529__aicore__ inline void603__aicore__ inline void
530-CSABlockVec<TEMPLATE_ARGS>::ProcessVec0(Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &outputL1,604+CSABlockVec<TEMPLATE_ARGS>::ProcessVec0(Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,
531- Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,605+ const RunInfo &runInfo, ConstInfo &constInfo)
532- const RunInfo &runInfo, ConstInfo &constInfo, int32_t startPos)
533{606{
534 if constexpr (TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE) {607 if constexpr (TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE) {
535 if (runInfo.s2LoopCount < runInfo.oriKvLoopEndIdx) {608 if (runInfo.s2LoopCount < runInfo.oriKvLoopEndIdx) {
536- if constexpr (IS_SPLIT_G) {
537- CrossCoreSetFlag<0, PIPE_MTE3>(15);
538- CrossCoreWaitFlag<0, PIPE_MTE3>(15);
539- }
540 return;609 return;
541 }610 }
542 keyGm = cmpKVGm;611 keyGm = cmpKVGm;
@@ -548,12 +617,7 @@ CSABlockVec<TEMPLATE_ARGS>::ProcessVec0(Buffer<BufferType::L1, SyncType::CROSS_C
548 maxBlockNumPerBatch = constInfo.cmpMaxBlockNumPerBatch;617 maxBlockNumPerBatch = constInfo.cmpMaxBlockNumPerBatch;
549 }618 }
550 CalSparseCalSize(runInfo, constInfo);619 CalSparseCalSize(runInfo, constInfo);
551- ProcessSparseKv(outputL1, v0ResGm, runInfo, constInfo, startPos);620+ ProcessSparseKv(v0ResGm, runInfo, constInfo);
552- if constexpr (IS_SPLIT_G) {
553- CrossCoreSetFlag<0, PIPE_MTE3>(15);
554- CrossCoreWaitFlag<0, PIPE_MTE3>(15);
555- }
556- outputL1.SetCrossCore();
557 v0ResGm.SetCrossCore();621 v0ResGm.SetCrossCore();
558 } else if constexpr (TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||622 } else if constexpr (TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||
559 TEMPLATE_MODE == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) {623 TEMPLATE_MODE == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) {
@@ -577,21 +641,15 @@ CSABlockVec<TEMPLATE_ARGS>::ProcessVec0(Buffer<BufferType::L1, SyncType::CROSS_C
577 }641 }
578 }642 }
579 CalSparseCalSize(runInfo, constInfo);643 CalSparseCalSize(runInfo, constInfo);
580- ProcessSparseKv(outputL1, v0ResGm, runInfo, constInfo, startPos);644+ ProcessSparseKv(v0ResGm, runInfo, constInfo);
581- if constexpr (IS_SPLIT_G) {
582- CrossCoreSetFlag<0, PIPE_MTE3>(15);
583- CrossCoreWaitFlag<0, PIPE_MTE3>(15);
584- }
585- outputL1.SetCrossCore();
586 v0ResGm.SetCrossCore();645 v0ResGm.SetCrossCore();
587 }646 }
588}647}
589 648 
590TEMPLATES_DEF_NO_DEFAULT649TEMPLATES_DEF_NO_DEFAULT
591__aicore__ inline void650__aicore__ inline void
592-CSABlockVec<TEMPLATE_ARGS>::ProcessSparseKv(Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &outputL1,651+CSABlockVec<TEMPLATE_ARGS>::ProcessSparseKv(Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,
593- Buffer<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> &v0ResGm,652+ const RunInfo &runInfo, ConstInfo &constInfo)
594- const RunInfo &runInfo, ConstInfo &constInfo, int32_t startPos)
595{653{
596 if constexpr (TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE ||654 if constexpr (TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE ||
597 TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||655 TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||
@@ -602,22 +660,26 @@ CSABlockVec<TEMPLATE_ARGS>::ProcessSparseKv(Buffer<BufferType::L1, SyncType::CRO
602 bool meetEnd = false;660 bool meetEnd = false;
603 int64_t s2Start = sparseS2Start;661 int64_t s2Start = sparseS2Start;
604 int64_t s2 = sparseS2Start;662 int64_t s2 = sparseS2Start;
605- int64_t token0Idx;
606- int64_t token1Idx;
607 uint32_t pingPong = 0;663 uint32_t pingPong = 0;
608 while ((s2 < sparseS2End) && !meetEnd) {664 while ((s2 < sparseS2End) && !meetEnd) {
609 int64_t dealRow = 0;665 int64_t dealRow = 0;
610 LocalTensor<Q_T> stage0OutUb = this->stage0OutBuf[pingPong].template Get<Q_T>();666 LocalTensor<Q_T> stage0OutUb = this->stage0OutBuf[pingPong].template Get<Q_T>();
611 WaitFlag<HardEvent::MTE3_MTE2>(mte3ToMte2[pingPong]);667 WaitFlag<HardEvent::MTE3_MTE2>(mte3ToMte2[pingPong]);
612 while (dealRow < Min(16, sparseCalSize) && s2 < sparseS2End) {668 while (dealRow < Min(16, sparseCalSize) && s2 < sparseS2End) {
613- GetRealCmpS2Idx(token0Idx, token1Idx, s2, runInfo, constInfo);669+ int64_t tokenData[8] = {-1, -1, -1, -1, -1, -1, -1, -1};
614- s2 += 2;670+ if constexpr (IS_VEC_S2PHYADDR) {
615- if (token0Idx == -1 && token1Idx == -1) {671+ GetRealS2Addr(tokenData, s2, runInfo, constInfo);
672+ } else {
673+ GetRealCmpS2Idx(tokenData, s2, runInfo, constInfo);
674+ }
675+ s2 += 8;
676+ if (tokenData[0] == -1 && tokenData[1] == -1 && tokenData[2] == -1 && tokenData[3] == -1 &&
677+ tokenData[4] == -1 && tokenData[5] == -1 && tokenData[6] == -1 && tokenData[7] == -1) {
616 meetEnd = true;678 meetEnd = true;
617 break;679 break;
618 }680 }
619- dealRow += CopyInKvSparse(stage0OutUb, dealRow, token0Idx, token1Idx, runInfo, constInfo);681+ dealRow += CopyInKvSparse(stage0OutUb, dealRow, tokenData, runInfo, constInfo);
620- if (token1Idx == -1) {682+ if (tokenData[7] == -1) {
621 meetEnd = true;683 meetEnd = true;
622 break;684 break;
623 }685 }
@@ -650,13 +712,13 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::ComputeVec1Softmax(
650 negativeFloatScalar);712 negativeFloatScalar);
651 } else if (runInfo.s2RealSize <= 64) {713 } else if (runInfo.s2RealSize <= 64) {
652 ProcessVec1Vf<T, Q_T, UPDATE, s1BaseSize, s2BaseSize,714 ProcessVec1Vf<T, Q_T, UPDATE, s1BaseSize, s2BaseSize,
653- FaVectorApi::OriginNRange::GT_0_AND_LTE_64_SFA>(715+ FaVectorApi::OriginNRange::GT_0_AND_LTE_64_SFA>(
654 stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, vselrIndexesBuf,716 stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, vselrIndexesBuf,
655 runInfo.halfMRealSize, runInfo.s2RealSizeUpdate, static_cast<T>(constInfo.softmaxScale),717 runInfo.halfMRealSize, runInfo.s2RealSizeUpdate, static_cast<T>(constInfo.softmaxScale),
656 negativeFloatScalar);718 negativeFloatScalar);
657 } else if (runInfo.s2RealSize < 128 || runInfo.s2RealSizeUpdate < 128) {719 } else if (runInfo.s2RealSize < 128 || runInfo.s2RealSizeUpdate < 128) {
658 ProcessVec1Vf<T, Q_T, UPDATE, s1BaseSize, s2BaseSize,720 ProcessVec1Vf<T, Q_T, UPDATE, s1BaseSize, s2BaseSize,
659- FaVectorApi::OriginNRange::GT_64_AND_LTE_128_SFA>(721+ FaVectorApi::OriginNRange::GT_64_AND_LTE_128_SFA>(
660 stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, vselrIndexesBuf,722 stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, vselrIndexesBuf,
661 runInfo.halfMRealSize, runInfo.s2RealSizeUpdate, static_cast<T>(constInfo.softmaxScale),723 runInfo.halfMRealSize, runInfo.s2RealSizeUpdate, static_cast<T>(constInfo.softmaxScale),
662 negativeFloatScalar);724 negativeFloatScalar);
@@ -700,10 +762,11 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::CopyVec1ResultToL1(
700 LocalTensor<Q_T> mm2AL1Tensor = outputBuf.GetTensor<Q_T>();762 LocalTensor<Q_T> mm2AL1Tensor = outputBuf.GetTensor<Q_T>();
701 if (likely(runInfo.halfMRealSize != 0)) {763 if (likely(runInfo.halfMRealSize != 0)) {
702 DataCopy(mm2AL1Tensor[constInfo.subBlockIdx * (BLOCK_BYTE / sizeof(Q_T)) *764 DataCopy(mm2AL1Tensor[constInfo.subBlockIdx * (BLOCK_BYTE / sizeof(Q_T)) *
703- (runInfo.mRealSize - runInfo.halfMRealSize)], stage1CastTensor,765+ (runInfo.mRealSize - runInfo.halfMRealSize)],
704- {s2BaseSize / 16, static_cast<uint16_t>(runInfo.halfMRealSize),766+ stage1CastTensor,
705- static_cast<uint16_t>(vec1Srcstride - runInfo.halfMRealSize),767+ {s2BaseSize / 16, static_cast<uint16_t>(runInfo.halfMRealSize),
706- static_cast<uint16_t>(Align16Func(runInfo.mRealSize) - runInfo.halfMRealSize)});768+ static_cast<uint16_t>(vec1Srcstride - runInfo.halfMRealSize),
769+ static_cast<uint16_t>(Align16Func(runInfo.mRealSize) - runInfo.halfMRealSize)});
707 }770 }
708 this->stage1OutQue[stage1Offset].template FreeTensor(stage1CastTensor);771 this->stage1OutQue[stage1Offset].template FreeTensor(stage1CastTensor);
709 outputBuf.SetCrossCore();772 outputBuf.SetCrossCore();
@@ -714,12 +777,12 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::StageCrossCoreVec1Lse(
714 LocalTensor<float> &maxUb, LocalTensor<float> &sumUb, RunInfo &runInfo, ConstInfo &constInfo)777 LocalTensor<float> &maxUb, LocalTensor<float> &sumUb, RunInfo &runInfo, ConstInfo &constInfo)
715{778{
716 AttentionCommon::S2SplitFdStagingLayout stagingLayout = {constInfo.gSize, dTemplateAlign64,779 AttentionCommon::S2SplitFdStagingLayout stagingLayout = {constInfo.gSize, dTemplateAlign64,
717- GetStagingSlotNum(false), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,780+ GetStagingSlotNum(false), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,
718- AttentionCommon::FD_REDUCE_CHUNK_ROWS};781+ AttentionCommon::FD_REDUCE_CHUNK_ROWS};
719 LocalTensor<float> tmpUb = this->batchReduceTmpBuf.template Get<float>();782 LocalTensor<float> tmpUb = this->batchReduceTmpBuf.template Get<float>();
720 AttentionCommon::StageVec1Lse(stagingLayout, crossCoreCombineBase, GetCrossCoreWorkspaceIdx(runInfo),783 AttentionCommon::StageVec1Lse(stagingLayout, crossCoreCombineBase, GetCrossCoreWorkspaceIdx(runInfo),
721- GetFaStagingMOffset(runInfo, constInfo), runInfo.halfMRealSize, maxUb, sumUb, tmpUb,784+ GetFaStagingMOffset(runInfo, constInfo), runInfo.halfMRealSize, maxUb, sumUb, tmpUb,
722- vToMte3AttnOutId, stageMte3ToVId);785+ vToMte3AttnOutId, stageMte3ToVId);
723}786}
724 787 
725TEMPLATES_DEF_NO_DEFAULT788TEMPLATES_DEF_NO_DEFAULT
@@ -741,14 +804,14 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::StageBatchConsistencyVec1Lse(
741 if (runInfo.isCrossCoreSplit && !runInfo.isFirstS2SplitCore) {804 if (runInfo.isCrossCoreSplit && !runInfo.isFirstS2SplitCore) {
742 StageCrossCoreVec1Lse(maxUb, sumUb, runInfo, constInfo);805 StageCrossCoreVec1Lse(maxUb, sumUb, runInfo, constInfo);
743 } else if (runInfo.isFirstS2SplitCore && runInfo.reduceBlockId == 0 &&806 } else if (runInfo.isFirstS2SplitCore && runInfo.reduceBlockId == 0 &&
744- runInfo.s2LoopCount < runInfo.s2LoopLimit) {807+ runInfo.s2LoopCount < runInfo.s2LoopLimit) {
745 AttentionCommon::S2SplitFdStagingLayout stagingLayout = {constInfo.gSize, dTemplateAlign64,808 AttentionCommon::S2SplitFdStagingLayout stagingLayout = {constInfo.gSize, dTemplateAlign64,
746- GetStagingSlotNum(true), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,809+ GetStagingSlotNum(true), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,
747- AttentionCommon::FD_REDUCE_CHUNK_ROWS};810+ AttentionCommon::FD_REDUCE_CHUNK_ROWS};
748 LocalTensor<float> tmpUb = this->batchReduceTmpBuf.template Get<float>();811 LocalTensor<float> tmpUb = this->batchReduceTmpBuf.template Get<float>();
749 AttentionCommon::StageVec1Lse(stagingLayout, intraCoreCombineBase,812 AttentionCommon::StageVec1Lse(stagingLayout, intraCoreCombineBase,
750- GetIntraCoreWorkspaceIdx(runInfo, constInfo), GetFaStagingMOffset(runInfo, constInfo),813+ GetIntraCoreWorkspaceIdx(runInfo, constInfo), GetFaStagingMOffset(runInfo, constInfo),
751- runInfo.halfMRealSize, maxUb, sumUb, tmpUb, vToMte3AttnOutId, stageMte3ToVId);814+ runInfo.halfMRealSize, maxUb, sumUb, tmpUb, vToMte3AttnOutId, stageMte3ToVId);
752 SetFlag<HardEvent::MTE3_MTE2>(intraLseMte3ToMte2Id[runInfo.multiCoreIdxMod2]);815 SetFlag<HardEvent::MTE3_MTE2>(intraLseMte3ToMte2Id[runInfo.multiCoreIdxMod2]);
753 } else if (runInfo.isCrossCoreSplit && runInfo.isFirstS2SplitCore && runInfo.reduceBlockId == 0) {816 } else if (runInfo.isCrossCoreSplit && runInfo.isFirstS2SplitCore && runInfo.reduceBlockId == 0) {
754 StageCrossCoreVec1Lse(maxUb, sumUb, runInfo, constInfo);817 StageCrossCoreVec1Lse(maxUb, sumUb, runInfo, constInfo);
@@ -764,12 +827,12 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::StageLegacyVec1Lse(
764 return;827 return;
765 }828 }
766 AttentionCommon::S2SplitFdStagingLayout stagingLayout = {constInfo.gSize, dTemplateAlign64,829 AttentionCommon::S2SplitFdStagingLayout stagingLayout = {constInfo.gSize, dTemplateAlign64,
767- GetStagingSlotNum(), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,830+ GetStagingSlotNum(), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,
768- AttentionCommon::FD_REDUCE_CHUNK_ROWS};831+ AttentionCommon::FD_REDUCE_CHUNK_ROWS};
769 LocalTensor<float> tmpUb = this->stage2OutBuf.template Get<float>();832 LocalTensor<float> tmpUb = this->stage2OutBuf.template Get<float>();
770 AttentionCommon::StageVec1Lse(stagingLayout, fdStagingBase, GetCrossCoreWorkspaceIdx(runInfo),833 AttentionCommon::StageVec1Lse(stagingLayout, fdStagingBase, GetCrossCoreWorkspaceIdx(runInfo),
771- GetFaStagingMOffset(runInfo, constInfo), static_cast<uint32_t>(runInfo.halfMRealSize),834+ GetFaStagingMOffset(runInfo, constInfo), static_cast<uint32_t>(runInfo.halfMRealSize),
772- maxUb, sumUb, tmpUb, vToMte3AttnOutId, stageMte3ToVId);835+ maxUb, sumUb, tmpUb, vToMte3AttnOutId, stageMte3ToVId);
773}836}
774 837 
775TEMPLATES_DEF_NO_DEFAULT838TEMPLATES_DEF_NO_DEFAULT
@@ -777,7 +840,7 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::CopyOutVec1Lse(
777 LocalTensor<float> &maxUb, LocalTensor<float> &sumUb, RunInfo &runInfo, ConstInfo &constInfo)840 LocalTensor<float> &maxUb, LocalTensor<float> &sumUb, RunInfo &runInfo, ConstInfo &constInfo)
778{841{
779 bool copyOutLse = constInfo.returnSoftmaxLse && runInfo.halfMRealSize > 0 &&842 bool copyOutLse = constInfo.returnSoftmaxLse && runInfo.halfMRealSize > 0 &&
780- runInfo.s2LoopCount == runInfo.s2LoopLimit;843+ runInfo.s2LoopCount == runInfo.s2LoopLimit;
781 if constexpr (IS_BATCH_CONSISTENCY) {844 if constexpr (IS_BATCH_CONSISTENCY) {
782 copyOutLse = copyOutLse && !runInfo.isCrossCoreSplit && !runInfo.needReduce;845 copyOutLse = copyOutLse && !runInfo.isCrossCoreSplit && !runInfo.needReduce;
783 }846 }
@@ -848,11 +911,11 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::ReduceIntraBlockAndStage(
848 RunInfo &runInfo, ConstInfo &constInfo, LocalTensor<T> &vec2ResUb, LocalTensor<T> &partialTmpUb)911 RunInfo &runInfo, ConstInfo &constInfo, LocalTensor<T> &vec2ResUb, LocalTensor<T> &partialTmpUb)
849{912{
850 AttentionCommon::S2SplitFdStagingLayout intraLayout = {constInfo.gSize, dTemplateAlign64,913 AttentionCommon::S2SplitFdStagingLayout intraLayout = {constInfo.gSize, dTemplateAlign64,
851- GetStagingSlotNum(true), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,914+ GetStagingSlotNum(true), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,
852- AttentionCommon::FD_REDUCE_CHUNK_ROWS};915+ AttentionCommon::FD_REDUCE_CHUNK_ROWS};
853 AttentionCommon::S2SplitFdStagingLayout crossLayout = {constInfo.gSize, dTemplateAlign64,916 AttentionCommon::S2SplitFdStagingLayout crossLayout = {constInfo.gSize, dTemplateAlign64,
854- GetStagingSlotNum(false), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,917+ GetStagingSlotNum(false), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,
855- AttentionCommon::FD_REDUCE_CHUNK_ROWS};918+ AttentionCommon::FD_REDUCE_CHUNK_ROWS};
856 uint32_t intraWorkspaceIdx = GetIntraCoreWorkspaceIdx(runInfo, constInfo);919 uint32_t intraWorkspaceIdx = GetIntraCoreWorkspaceIdx(runInfo, constInfo);
857 uint32_t crossWorkspaceIdx = static_cast<uint32_t>(920 uint32_t crossWorkspaceIdx = static_cast<uint32_t>(
858 runInfo.firstFdDataWorkspaceIdx + runInfo.s2SplitIdx - runInfo.reduceBlockId);921 runInfo.firstFdDataWorkspaceIdx + runInfo.s2SplitIdx - runInfo.reduceBlockId);
@@ -867,7 +930,7 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::ReduceIntraBlockAndStage(
867 LocalTensor<float> sumUb =930 LocalTensor<float> sumUb =
868 this->softmaxFinalSumBuf[runInfo.taskIdMod2].template Get<float>();931 this->softmaxFinalSumBuf[runInfo.taskIdMod2].template Get<float>();
869 bool copyOutMergedLse = constInfo.returnSoftmaxLse && !runInfo.isCrossCoreSplit &&932 bool copyOutMergedLse = constInfo.returnSoftmaxLse && !runInfo.isCrossCoreSplit &&
870- runInfo.s2LoopCount == runInfo.s2LoopLimit;933+ runInfo.s2LoopCount == runInfo.s2LoopLimit;
871 934 
872 WaitFlag<HardEvent::MTE3_MTE2>(intraLseMte3ToMte2Id[runInfo.multiCoreIdxMod2]);935 WaitFlag<HardEvent::MTE3_MTE2>(intraLseMte3ToMte2Id[runInfo.multiCoreIdxMod2]);
873 WaitFlag<HardEvent::MTE3_MTE2>(intraAttnOutMte3ToMte2Id[runInfo.multiCoreIdxMod2]);936 WaitFlag<HardEvent::MTE3_MTE2>(intraAttnOutMte3ToMte2Id[runInfo.multiCoreIdxMod2]);
@@ -884,14 +947,14 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::ReduceIntraBlockAndStage(
884 WaitFlag<HardEvent::MTE3_V>(mte3ToVLseOutId);947 WaitFlag<HardEvent::MTE3_V>(mte3ToVLseOutId);
885 }948 }
886 AttentionCommon::MergeStagedAndCurrentChunk<T, dTemplateAlign64>(intraLayout,949 AttentionCommon::MergeStagedAndCurrentChunk<T, dTemplateAlign64>(intraLayout,
887- intraCoreCombineBase, intraWorkspaceIdx, stagingMOffset + startRow,950+ intraCoreCombineBase, intraWorkspaceIdx, stagingMOffset + startRow,
888- dealRowCount, static_cast<int64_t>(constInfo.dSizeV), chunkMaxUb, chunkSumUb,951+ dealRowCount, static_cast<int64_t>(constInfo.dSizeV), chunkMaxUb, chunkSumUb,
889- chunkCurrent, blockMaxUb, blockSumUb, partialTmpUb, lseBroadcastUb, sumBroadcastUb,952+ chunkCurrent, blockMaxUb, blockSumUb, partialTmpUb, lseBroadcastUb, sumBroadcastUb,
890- reduceMaxSumVToMte2Id, intraPartialOVToMte2Id, reduceMte2ToVId);953+ reduceMaxSumVToMte2Id, intraPartialOVToMte2Id, reduceMte2ToVId);
891 954 
892 AttentionCommon::StageBroadcastMaxSum(intraLayout, intraCoreCombineBase, intraWorkspaceIdx,955 AttentionCommon::StageBroadcastMaxSum(intraLayout, intraCoreCombineBase, intraWorkspaceIdx,
893- stagingMOffset + startRow, dealRowCount, lseBroadcastUb, sumBroadcastUb,956+ stagingMOffset + startRow, dealRowCount, lseBroadcastUb, sumBroadcastUb,
894- vToMte3AttnOutId, stageMte3ToVId);957+ vToMte3AttnOutId, stageMte3ToVId);
895 if (copyOutMergedLse) {958 if (copyOutMergedLse) {
896 DataCopyExtParams lseParams;959 DataCopyExtParams lseParams;
897 lseParams.blockCount = static_cast<uint16_t>(dealRowCount);960 lseParams.blockCount = static_cast<uint16_t>(dealRowCount);
@@ -901,24 +964,24 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::ReduceIntraBlockAndStage(
901 SetFlag<HardEvent::V_MTE3>(vToMte3LseOutId);964 SetFlag<HardEvent::V_MTE3>(vToMte3LseOutId);
902 WaitFlag<HardEvent::V_MTE3>(vToMte3LseOutId);965 WaitFlag<HardEvent::V_MTE3>(vToMte3LseOutId);
903 DataCopyPad(this->softmaxLseGm[runInfo.softmaxLseOffset + startRow],966 DataCopyPad(this->softmaxLseGm[runInfo.softmaxLseOffset + startRow],
904- lseBroadcastUb, lseParams);967+ lseBroadcastUb, lseParams);
905 SetFlag<HardEvent::MTE3_V>(mte3ToVLseOutId);968 SetFlag<HardEvent::MTE3_V>(mte3ToVLseOutId);
906 }969 }
907 if (runInfo.isCrossCoreSplit && runInfo.s2LoopCount == runInfo.s2LoopLimit) {970 if (runInfo.isCrossCoreSplit && runInfo.s2LoopCount == runInfo.s2LoopLimit) {
908 AttentionCommon::StageBroadcastMaxSum(crossLayout, crossCoreCombineBase, crossWorkspaceIdx,971 AttentionCommon::StageBroadcastMaxSum(crossLayout, crossCoreCombineBase, crossWorkspaceIdx,
909- stagingMOffset + startRow, dealRowCount, lseBroadcastUb, sumBroadcastUb,972+ stagingMOffset + startRow, dealRowCount, lseBroadcastUb, sumBroadcastUb,
910- vToMte3AttnOutId, stageMte3ToVId);973+ vToMte3AttnOutId, stageMte3ToVId);
911 }974 }
912 startRow += intraLayout.chunkRows;975 startRow += intraLayout.chunkRows;
913 }976 }
914 977 
915 AttentionCommon::StageVec2PartialOAndWait<T>(intraLayout, intraCoreCombineGm, intraWorkspaceIdx,978 AttentionCommon::StageVec2PartialOAndWait<T>(intraLayout, intraCoreCombineGm, intraWorkspaceIdx,
916- stagingMOffset, runInfo.vec2MRealSize, static_cast<uint32_t>(constInfo.dSizeV), vec2ResUb,979+ stagingMOffset, runInfo.vec2MRealSize, static_cast<uint32_t>(constInfo.dSizeV), vec2ResUb,
917- vToMte3AttnOutId, stageMte3ToVId);980+ vToMte3AttnOutId, stageMte3ToVId);
918 if (runInfo.isCrossCoreSplit && runInfo.s2LoopCount == runInfo.s2LoopLimit) {981 if (runInfo.isCrossCoreSplit && runInfo.s2LoopCount == runInfo.s2LoopLimit) {
919 AttentionCommon::StageVec2PartialOAndWait<T>(crossLayout, crossCoreCombineGm, crossWorkspaceIdx,982 AttentionCommon::StageVec2PartialOAndWait<T>(crossLayout, crossCoreCombineGm, crossWorkspaceIdx,
920- stagingMOffset, runInfo.vec2MRealSize, static_cast<uint32_t>(constInfo.dSizeV), vec2ResUb,983+ stagingMOffset, runInfo.vec2MRealSize, static_cast<uint32_t>(constInfo.dSizeV), vec2ResUb,
921- vToMte3AttnOutId, stageMte3ToVId);984+ vToMte3AttnOutId, stageMte3ToVId);
922 }985 }
923 if (runInfo.s2LoopCount < runInfo.s2LoopLimit) {986 if (runInfo.s2LoopCount < runInfo.s2LoopLimit) {
924 SetFlag<HardEvent::MTE3_MTE2>(intraLseMte3ToMte2Id[runInfo.multiCoreIdxMod2]);987 SetFlag<HardEvent::MTE3_MTE2>(intraLseMte3ToMte2Id[runInfo.multiCoreIdxMod2]);
@@ -1016,27 +1079,27 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::ProcessVec2(
1016 if (runInfo.isFirstS2SplitCore && runInfo.reduceBlockId == 0 &&1079 if (runInfo.isFirstS2SplitCore && runInfo.reduceBlockId == 0 &&
1017 runInfo.s2LoopCount < runInfo.s2LoopLimit) {1080 runInfo.s2LoopCount < runInfo.s2LoopLimit) {
1018 AttentionCommon::S2SplitFdStagingLayout stagingLayout = {constInfo.gSize,1081 AttentionCommon::S2SplitFdStagingLayout stagingLayout = {constInfo.gSize,
1019- dTemplateAlign64, GetStagingSlotNum(true), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,1082+ dTemplateAlign64, GetStagingSlotNum(true), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,
1020- AttentionCommon::FD_REDUCE_CHUNK_ROWS};1083+ AttentionCommon::FD_REDUCE_CHUNK_ROWS};
1021 int64_t stagingMOffset = GetFaStagingMOffset(runInfo, constInfo);1084 int64_t stagingMOffset = GetFaStagingMOffset(runInfo, constInfo);
1022 AttentionCommon::StageVec2PartialOAndWait<T>(stagingLayout, intraCoreCombineGm,1085 AttentionCommon::StageVec2PartialOAndWait<T>(stagingLayout, intraCoreCombineGm,
1023- GetIntraCoreWorkspaceIdx(runInfo, constInfo), stagingMOffset, runInfo.vec2MRealSize,1086+ GetIntraCoreWorkspaceIdx(runInfo, constInfo), stagingMOffset, runInfo.vec2MRealSize,
1024- static_cast<uint32_t>(constInfo.dSizeV), vec2ResUb,1087+ static_cast<uint32_t>(constInfo.dSizeV), vec2ResUb,
1025- vToMte3AttnOutId, stageMte3ToVId);1088+ vToMte3AttnOutId, stageMte3ToVId);
1026 SetFlag<HardEvent::MTE3_MTE2>(intraAttnOutMte3ToMte2Id[runInfo.multiCoreIdxMod2]);1089 SetFlag<HardEvent::MTE3_MTE2>(intraAttnOutMte3ToMte2Id[runInfo.multiCoreIdxMod2]);
1027 }1090 }
1028 if (runInfo.isCrossCoreSplit && (!runInfo.isFirstS2SplitCore ||1091 if (runInfo.isCrossCoreSplit && (!runInfo.isFirstS2SplitCore ||
1029- (runInfo.isFirstS2SplitCore && runInfo.reduceBlockId == 0 &&1092+ (runInfo.isFirstS2SplitCore && runInfo.reduceBlockId == 0 &&
1030- runInfo.s2LoopCount == runInfo.s2LoopLimit))) {1093+ runInfo.s2LoopCount == runInfo.s2LoopLimit))) {
1031 AttentionCommon::S2SplitFdStagingLayout stagingLayout = {constInfo.gSize,1094 AttentionCommon::S2SplitFdStagingLayout stagingLayout = {constInfo.gSize,
1032- dTemplateAlign64, GetStagingSlotNum(false), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,1095+ dTemplateAlign64, GetStagingSlotNum(false), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,
1033- AttentionCommon::FD_REDUCE_CHUNK_ROWS};1096+ AttentionCommon::FD_REDUCE_CHUNK_ROWS};
1034 uint32_t workspaceIdx = GetCrossCoreWorkspaceIdx(runInfo);1097 uint32_t workspaceIdx = GetCrossCoreWorkspaceIdx(runInfo);
1035 int64_t stagingMOffset = GetFaStagingMOffset(runInfo, constInfo);1098 int64_t stagingMOffset = GetFaStagingMOffset(runInfo, constInfo);
1036 AttentionCommon::StageVec2PartialOAndWait<T>(stagingLayout, crossCoreCombineGm,1099 AttentionCommon::StageVec2PartialOAndWait<T>(stagingLayout, crossCoreCombineGm,
1037- workspaceIdx, stagingMOffset, runInfo.vec2MRealSize,1100+ workspaceIdx, stagingMOffset, runInfo.vec2MRealSize,
1038- static_cast<uint32_t>(constInfo.dSizeV), vec2ResUb,1101+ static_cast<uint32_t>(constInfo.dSizeV), vec2ResUb,
1039- vToMte3AttnOutId, stageMte3ToVId);1102+ vToMte3AttnOutId, stageMte3ToVId);
1040 } else if (!runInfo.isCrossCoreSplit && runInfo.s2LoopCount == runInfo.s2LoopLimit) {1103 } else if (!runInfo.isCrossCoreSplit && runInfo.s2LoopCount == runInfo.s2LoopLimit) {
1041 this->CopyOutAttentionOut(runInfo, constInfo, vec2ResUb, 0, vec2CalcSize);1104 this->CopyOutAttentionOut(runInfo, constInfo, vec2ResUb, 0, vec2CalcSize);
1042 }1105 }
@@ -1086,14 +1149,15 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::ProcessFlashDecode(FdRunInfo
1086 seqOffset = fdRunInfo.bn2Idx * constInfo.s1Size;1149 seqOffset = fdRunInfo.bn2Idx * constInfo.s1Size;
1087 }1150 }
1088 int64_t attentionOutOffset = seqOffset * constInfo.n2GDv +1151 int64_t attentionOutOffset = seqOffset * constInfo.n2GDv +
1089- fdRunInfo.mIdx * constInfo.n2GDv +1152+ fdRunInfo.mIdx * constInfo.n2GDv +
1090- fdRunInfo.mStartIdx * constInfo.dSizeV;1153+ fdRunInfo.mStartIdx * constInfo.dSizeV;
1091 int64_t softmaxLseOffset = 0;1154 int64_t softmaxLseOffset = 0;
1092 if constexpr (LAYOUT_T == SMLA_LAYOUT::TND) {1155 if constexpr (LAYOUT_T == SMLA_LAYOUT::TND) {
1093 softmaxLseOffset = (seqOffset + fdRunInfo.mIdx) * constInfo.gSize + fdRunInfo.mStartIdx;1156 softmaxLseOffset = (seqOffset + fdRunInfo.mIdx) * constInfo.gSize + fdRunInfo.mStartIdx;
1094 } else {1157 } else {
1095 softmaxLseOffset = (fdRunInfo.bn2Idx * constInfo.s1Size + fdRunInfo.mIdx) *1158 softmaxLseOffset = (fdRunInfo.bn2Idx * constInfo.s1Size + fdRunInfo.mIdx) *
1096- constInfo.gSize + fdRunInfo.mStartIdx;1159+ constInfo.gSize +
1160+ fdRunInfo.mStartIdx;
1097 }1161 }
1098 LocalTensor<T> accumulatedO = this->fdBuffers.accumOut.template Get<T>();1162 LocalTensor<T> accumulatedO = this->fdBuffers.accumOut.template Get<T>();
1099 LocalTensor<float> lseExpUb = this->fdBuffers.lseExp.template Get<float>();1163 LocalTensor<float> lseExpUb = this->fdBuffers.lseExp.template Get<float>();
@@ -1101,10 +1165,10 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::ProcessFlashDecode(FdRunInfo
1101 LocalTensor<float> blockSumUb = this->fdBuffers.blockSum.template Get<float>();1165 LocalTensor<float> blockSumUb = this->fdBuffers.blockSum.template Get<float>();
1102 LocalTensor<T> partialOFp32 = this->fdBuffers.partialO.template Get<T>();1166 LocalTensor<T> partialOFp32 = this->fdBuffers.partialO.template Get<T>();
1103 AttentionCommon::S2SplitFdStagingLayout stagingLayout = {constInfo.gSize,1167 AttentionCommon::S2SplitFdStagingLayout stagingLayout = {constInfo.gSize,
1104- dTemplateAlign64, GetStagingSlotNum(), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,1168+ dTemplateAlign64, GetStagingSlotNum(), AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW,
1105- AttentionCommon::FD_REDUCE_CHUNK_ROWS};1169+ AttentionCommon::FD_REDUCE_CHUNK_ROWS};
1106 int64_t attentionOutRowStride = static_cast<int64_t>(constInfo.dSizeV) +1170 int64_t attentionOutRowStride = static_cast<int64_t>(constInfo.dSizeV) +
1107- static_cast<int64_t>(constInfo.attentionOutStride) / sizeof(OUTPUT_T);1171+ static_cast<int64_t>(constInfo.attentionOutStride) / sizeof(OUTPUT_T);
1108 int64_t startRow = 0;1172 int64_t startRow = 0;
1109 while (startRow < fdRunInfo.mNum) {1173 while (startRow < fdRunInfo.mNum) {
1110 int64_t dealRowCount = AttentionCommon::FD_REDUCE_CHUNK_ROWS;1174 int64_t dealRowCount = AttentionCommon::FD_REDUCE_CHUNK_ROWS;
@@ -1113,13 +1177,13 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::ProcessFlashDecode(FdRunInfo
1113 }1177 }
1114 WaitFlag<HardEvent::MTE3_V>(fdMte3ToVId);1178 WaitFlag<HardEvent::MTE3_V>(fdMte3ToVId);
1115 AttentionCommon::ReduceWithLse<T, dTemplateAlign64>(stagingLayout, fdStagingBase,1179 AttentionCommon::ReduceWithLse<T, dTemplateAlign64>(stagingLayout, fdStagingBase,
1116- fdRunInfo.workspaceIdx, fdRunInfo.workspaceNum,1180+ fdRunInfo.workspaceIdx, fdRunInfo.workspaceNum,
1117- static_cast<uint32_t>(fdRunInfo.mStartIdx + startRow), dealRowCount,1181+ static_cast<uint32_t>(fdRunInfo.mStartIdx + startRow), dealRowCount,
1118- static_cast<uint32_t>(constInfo.dSizeV),1182+ static_cast<uint32_t>(constInfo.dSizeV),
1119- accumulatedO, lseExpUb, blockMaxUb, blockSumUb, partialOFp32,1183+ accumulatedO, lseExpUb, blockMaxUb, blockSumUb, partialOFp32,
1120- constInfo.returnSoftmaxLse, softmaxLseGm, softmaxLseOffset + startRow,1184+ constInfo.returnSoftmaxLse, softmaxLseGm, softmaxLseOffset + startRow,
1121- fdVToMte2Id[0], fdVToMte2Id[1], fdMte2ToVId,1185+ fdVToMte2Id[0], fdVToMte2Id[1], fdMte2ToVId,
1122- vToMte3LseOutId, mte3ToVLseOutId);1186+ vToMte3LseOutId, mte3ToVLseOutId);
1123 RunInfo runInfo;1187 RunInfo runInfo;
1124 runInfo.vec2MRealSize = dealRowCount;1188 runInfo.vec2MRealSize = dealRowCount;
1125 runInfo.attentionOutOffset = attentionOutOffset + startRow * attentionOutRowStride;1189 runInfo.attentionOutOffset = attentionOutOffset + startRow * attentionOutRowStride;
@@ -1373,6 +1437,773 @@ __aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::GetExtremeValue(T &negativeSc
1373 negativeScalar = *((float *)&tmp1);1437 negativeScalar = *((float *)&tmp1);
1374}1438}
1375 1439 
1440+template <typename T>
1441+__simd_vf__ void GetKVPhyAddrVFPaImpl(__ubuf__ uint32_t *kvPhyAddrUb, __ubuf__ int32_t *sparseIdxUb,
1442+ __ubuf__ int32_t *blkTableUb, const uint16_t s2Loop, uint32_t s2Tail,
1443+ const uint32_t blockSize, const int16_t shiftRightNum,
1444+ const uint32_t sparseBlockSize, const uint32_t kvDim, const uint32_t kvStride)
1445+{
1446+ static const uint16_t s2_num_per_loop = 128;
1447+ static const uint16_t s2_num_per_reg = 64;
1448+ static const uint16_t out_offset_per_loop = 256;
1449+ static const uint16_t out_offset_per_reg = 128;
1450+ static const uint32_t invalid_value = 0xFFFFFFFF;
1451+ MicroAPI::MaskReg preg_all_b32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
1452+ MicroAPI::MaskReg add_carry_l_1;
1453+ MicroAPI::MaskReg add_carry_h_1;
1454+ MicroAPI::MaskReg add_carry_l_2;
1455+ MicroAPI::MaskReg add_carry_h_2;
1456+ MicroAPI::MaskReg preg_tail_neg_1_b32;
1457+ MicroAPI::MaskReg preg_tail_neg_2_b32;
1458+ 
1459+ MicroAPI::RegTensor<uint32_t> vreg_kv_stride;
1460+ MicroAPI::RegTensor<uint32_t> vreg_sparse_idx_1;
1461+ MicroAPI::RegTensor<uint32_t> vreg_sparse_idx_2;
1462+ MicroAPI::RegTensor<uint32_t> vreg_block_size;
1463+ MicroAPI::RegTensor<uint32_t> vreg_shift_rights_num;
1464+ MicroAPI::RegTensor<uint32_t> vreg_pa_blk_idx_1;
1465+ MicroAPI::RegTensor<uint32_t> vreg_pa_blk_idx_2;
1466+ MicroAPI::RegTensor<uint32_t> vreg_pa_tmp_1;
1467+ MicroAPI::RegTensor<uint32_t> vreg_pa_tmp_2;
1468+ MicroAPI::RegTensor<uint32_t> vreg_pa_offset_1;
1469+ MicroAPI::RegTensor<uint32_t> vreg_pa_offset_2;
1470+ MicroAPI::RegTensor<uint32_t> vreg_phy_offset_1;
1471+ MicroAPI::RegTensor<uint32_t> vreg_phy_offset_2;
1472+ MicroAPI::RegTensor<uint32_t> vreg_phy_blk_idx_1;
1473+ MicroAPI::RegTensor<uint32_t> vreg_phy_blk_idx_2;
1474+ 
1475+ MicroAPI::RegTensor<uint32_t> vreg_blk_id_mul_stride_h_1;
1476+ MicroAPI::RegTensor<uint32_t> vreg_blk_id_mul_stride_tmp_h_1;
1477+ MicroAPI::RegTensor<uint32_t> vreg_blk_id_mul_stride_l_1;
1478+ MicroAPI::RegTensor<uint32_t> vreg_mul_overflow_l_1;
1479+ MicroAPI::RegTensor<uint32_t> vreg_total_offset_l_1;
1480+ MicroAPI::RegTensor<uint32_t> vreg_total_offset_h_1;
1481+ 
1482+ MicroAPI::RegTensor<uint32_t> vreg_blk_id_mul_stride_h_2;
1483+ MicroAPI::RegTensor<uint32_t> vreg_blk_id_mul_stride_tmp_h_2;
1484+ MicroAPI::RegTensor<uint32_t> vreg_blk_id_mul_stride_l_2;
1485+ MicroAPI::RegTensor<uint32_t> vreg_mul_overflow_l_2;
1486+ MicroAPI::RegTensor<uint32_t> vreg_total_offset_l_2;
1487+ MicroAPI::RegTensor<uint32_t> vreg_total_offset_h_2;
1488+ 
1489+ MicroAPI::RegTensor<uint32_t> vreg_zero;
1490+ MicroAPI::Duplicate(vreg_zero, 0);
1491+ MicroAPI::Duplicate(vreg_kv_stride, kvStride);
1492+ 
1493+ for (; s2Loop > 1;) {
1494+ for (uint16_t i = 0; i < s2Loop - 1; i++) {
1495+ MicroAPI::LoadAlign<int32_t, MicroAPI::LoadDist::DIST_NORM>(
1496+ (MicroAPI::RegTensor<int32_t> &)vreg_sparse_idx_1, sparseIdxUb + i * s2_num_per_loop);
1497+ MicroAPI::LoadAlign<int32_t, MicroAPI::LoadDist::DIST_NORM>(
1498+ (MicroAPI::RegTensor<int32_t> &)vreg_sparse_idx_2, sparseIdxUb + s2_num_per_reg + i * s2_num_per_loop);
1499+ // * sparseBlockSize
1500+ MicroAPI::Muls(vreg_sparse_idx_1, vreg_sparse_idx_1, sparseBlockSize, preg_all_b32);
1501+ MicroAPI::Muls(vreg_sparse_idx_2, vreg_sparse_idx_2, sparseBlockSize, preg_all_b32);
1502+ // 计算右移位数
1503+ // 右移 -> 除blockSize 得到paBlockIdx，vreg_sparse_idx - pa_idx * blocksize -> pa offset
1504+ MicroAPI::ShiftRights(vreg_pa_blk_idx_1, vreg_sparse_idx_1, shiftRightNum, preg_all_b32);
1505+ MicroAPI::ShiftRights(vreg_pa_blk_idx_2, vreg_sparse_idx_2, shiftRightNum, preg_all_b32);
1506+ 
1507+ MicroAPI::Muls(vreg_pa_tmp_1, vreg_pa_blk_idx_1, blockSize, preg_all_b32);
1508+ MicroAPI::Muls(vreg_pa_tmp_2, vreg_pa_blk_idx_2, blockSize, preg_all_b32);
1509+ // offset
1510+ MicroAPI::Sub(vreg_pa_offset_1, vreg_sparse_idx_1, vreg_pa_tmp_1, preg_all_b32);
1511+ MicroAPI::Sub(vreg_pa_offset_2, vreg_sparse_idx_2, vreg_pa_tmp_2, preg_all_b32);
1512+ // 物理页内offset
1513+ MicroAPI::Muls(vreg_phy_offset_1, vreg_pa_offset_1, kvDim, preg_all_b32);
1514+ MicroAPI::Muls(vreg_phy_offset_2, vreg_pa_offset_2, kvDim, preg_all_b32);
1515+ 
1516+ // int32 paBlockId -> 物理id
1517+ DataCopyGather(vreg_phy_blk_idx_1, blkTableUb, vreg_pa_blk_idx_1, preg_all_b32);
1518+ DataCopyGather(vreg_phy_blk_idx_2, blkTableUb, vreg_pa_blk_idx_2, preg_all_b32);
1519+ 
1520+ // 分高低32位计算int64物理地址 -- 乘 stride
1521+ // 低位乘 带进位
1522+ MicroAPI::Mull(vreg_blk_id_mul_stride_l_1, vreg_mul_overflow_l_1, vreg_phy_blk_idx_1, vreg_kv_stride,
1523+ preg_all_b32);
1524+ MicroAPI::Mull(vreg_blk_id_mul_stride_l_2, vreg_mul_overflow_l_2, vreg_phy_blk_idx_2, vreg_kv_stride,
1525+ preg_all_b32);
1526+ 
1527+ // 分高低32位计算int64物理地址 -- 加 offset
1528+ MicroAPI::Add(add_carry_l_1, vreg_total_offset_l_1, vreg_blk_id_mul_stride_l_1, vreg_phy_offset_1,
1529+ preg_all_b32);
1530+ MicroAPI::Add(add_carry_l_2, vreg_total_offset_l_2, vreg_blk_id_mul_stride_l_2, vreg_phy_offset_2,
1531+ preg_all_b32);
1532+ 
1533+ MicroAPI::AddC(add_carry_h_1, vreg_total_offset_h_1, vreg_mul_overflow_l_1, vreg_zero, add_carry_l_1,
1534+ preg_all_b32);
1535+ MicroAPI::AddC(add_carry_h_2, vreg_total_offset_h_2, vreg_mul_overflow_l_2, vreg_zero, add_carry_l_2,
1536+ preg_all_b32);
1537+ 
1538+ // 搬出 由于拆分为了int32类型，元素个数翻倍
1539+ MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(
1540+ kvPhyAddrUb + i * out_offset_per_loop, vreg_total_offset_l_1, vreg_total_offset_h_1, preg_all_b32);
1541+ MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(
1542+ kvPhyAddrUb + out_offset_per_reg + i * out_offset_per_loop, vreg_total_offset_l_2,
1543+ vreg_total_offset_h_2, preg_all_b32);
1544+ }
1545+ break;
1546+ }
1547+ 
1548+ for (uint16_t i = s2Loop - 1; i < s2Loop; i++) {
1549+ MicroAPI::MaskReg preg_tail_1_b32 = MicroAPI::UpdateMask<int32_t>(s2Tail);
1550+ MicroAPI::MaskReg preg_tail_2_b32 = MicroAPI::UpdateMask<int32_t>(s2Tail);
1551+ MicroAPI::Not(preg_tail_neg_1_b32, preg_tail_1_b32, preg_all_b32);
1552+ MicroAPI::Not(preg_tail_neg_2_b32, preg_tail_2_b32, preg_all_b32);
1553+ 
1554+ MicroAPI::LoadAlign<int32_t, MicroAPI::LoadDist::DIST_NORM>((MicroAPI::RegTensor<int32_t> &)vreg_sparse_idx_1,
1555+ sparseIdxUb + i * s2_num_per_loop);
1556+ MicroAPI::LoadAlign<int32_t, MicroAPI::LoadDist::DIST_NORM>((MicroAPI::RegTensor<int32_t> &)vreg_sparse_idx_2,
1557+ sparseIdxUb + s2_num_per_reg + i * s2_num_per_loop);
1558+ // * sparseBlockSize
1559+ MicroAPI::Muls(vreg_sparse_idx_1, vreg_sparse_idx_1, sparseBlockSize, preg_tail_1_b32);
1560+ MicroAPI::Muls(vreg_sparse_idx_2, vreg_sparse_idx_2, sparseBlockSize, preg_tail_2_b32);
1561+ // 计算右移位数
1562+ // 右移 -> 除blockSize 得到paBlockIdx，vreg_sparse_idx - pa_idx * blocksize -> pa offset
1563+ MicroAPI::ShiftRights(vreg_pa_blk_idx_1, vreg_sparse_idx_1, shiftRightNum, preg_tail_1_b32);
1564+ MicroAPI::ShiftRights(vreg_pa_blk_idx_2, vreg_sparse_idx_2, shiftRightNum, preg_tail_2_b32);
1565+ 
1566+ MicroAPI::Muls(vreg_pa_tmp_1, vreg_pa_blk_idx_1, blockSize, preg_tail_1_b32);
1567+ MicroAPI::Muls(vreg_pa_tmp_2, vreg_pa_blk_idx_2, blockSize, preg_tail_2_b32);
1568+ // offset
1569+ MicroAPI::Sub(vreg_pa_offset_1, vreg_sparse_idx_1, vreg_pa_tmp_1, preg_tail_1_b32);
1570+ MicroAPI::Sub(vreg_pa_offset_2, vreg_sparse_idx_2, vreg_pa_tmp_2, preg_tail_2_b32);
1571+ // 物理页内offset
1572+ MicroAPI::Muls(vreg_phy_offset_1, vreg_pa_offset_1, kvDim, preg_tail_1_b32);
1573+ MicroAPI::Muls(vreg_phy_offset_2, vreg_pa_offset_2, kvDim, preg_tail_2_b32);
1574+ 
1575+ // int32 paBlockId -> 物理id
1576+ DataCopyGather(vreg_phy_blk_idx_1, blkTableUb, vreg_pa_blk_idx_1, preg_tail_1_b32);
1577+ DataCopyGather(vreg_phy_blk_idx_2, blkTableUb, vreg_pa_blk_idx_2, preg_tail_2_b32);
1578+ 
1579+ // 分高低32位计算int64物理地址 -- 乘 stride
1580+ // 低位乘 带进位
1581+ MicroAPI::Mull(vreg_blk_id_mul_stride_l_1, vreg_mul_overflow_l_1, vreg_phy_blk_idx_1, vreg_kv_stride,
1582+ preg_tail_1_b32);
1583+ MicroAPI::Mull(vreg_blk_id_mul_stride_l_2, vreg_mul_overflow_l_2, vreg_phy_blk_idx_2, vreg_kv_stride,
1584+ preg_tail_2_b32);
1585+ 
1586+ // 分高低32位计算int64物理地址 -- 加 offset
1587+ MicroAPI::Add(add_carry_l_1, vreg_total_offset_l_1, vreg_blk_id_mul_stride_l_1, vreg_phy_offset_1,
1588+ preg_tail_1_b32);
1589+ MicroAPI::Add(add_carry_l_2, vreg_total_offset_l_2, vreg_blk_id_mul_stride_l_2, vreg_phy_offset_2,
1590+ preg_tail_2_b32);
1591+ 
1592+ MicroAPI::AddC(add_carry_h_1, vreg_total_offset_h_1, vreg_mul_overflow_l_1, vreg_zero, add_carry_l_1,
1593+ preg_tail_1_b32);
1594+ MicroAPI::AddC(add_carry_h_2, vreg_total_offset_h_2, vreg_mul_overflow_l_2, vreg_zero, add_carry_l_2,
1595+ preg_tail_2_b32);
1596+ 
1597+ // 无效值填充-1(0xFFFFFFFF)
1598+ MicroAPI::Duplicate<uint32_t, MicroAPI::MaskMergeMode::MERGING>(vreg_total_offset_l_1, invalid_value,
1599+ preg_tail_neg_1_b32);
1600+ MicroAPI::Duplicate<uint32_t, MicroAPI::MaskMergeMode::MERGING>(vreg_total_offset_h_1, invalid_value,
1601+ preg_tail_neg_1_b32);
1602+ MicroAPI::Duplicate<uint32_t, MicroAPI::MaskMergeMode::MERGING>(vreg_total_offset_l_2, invalid_value,
1603+ preg_tail_neg_2_b32);
1604+ MicroAPI::Duplicate<uint32_t, MicroAPI::MaskMergeMode::MERGING>(vreg_total_offset_h_2, invalid_value,
1605+ preg_tail_neg_2_b32);
1606+ MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(
1607+ kvPhyAddrUb + i * out_offset_per_loop, vreg_total_offset_l_1, vreg_total_offset_h_1, preg_all_b32);
1608+ MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(
1609+ kvPhyAddrUb + out_offset_per_reg + i * out_offset_per_loop, vreg_total_offset_l_2,
1610+ vreg_total_offset_h_2, preg_all_b32);
1611+ }
1612+}
1613+ 
1614+template <typename T>
1615+__aicore__ inline void GetKVPhyAddrVFPa(LocalTensor<uint32_t> kvPhyAddrTensor, LocalTensor<int32_t> sparseIdxTensor,
1616+ LocalTensor<int32_t> blkTableTensor, const uint16_t s2Loop, const uint32_t s2Tail,
1617+ const uint32_t blockSize, const int16_t shiftRightNum,
1618+ const uint32_t sparseBlockSize, const uint32_t kvDim, const uint32_t kvStride)
1619+{
1620+ __ubuf__ uint32_t *kv_phy_addr_ub = (__ubuf__ uint32_t *)(kvPhyAddrTensor.GetPhyAddr());
1621+ __ubuf__ int32_t *sparse_idx_ub = (__ubuf__ int32_t *)(sparseIdxTensor.GetPhyAddr());
1622+ __ubuf__ int32_t *blk_table_ub = (__ubuf__ int32_t *)(blkTableTensor.GetPhyAddr());
1623+ GetKVPhyAddrVFPaImpl<uint32_t>(kv_phy_addr_ub, sparse_idx_ub, blk_table_ub, s2Loop, s2Tail, blockSize, shiftRightNum,
1624+ sparseBlockSize, kvDim, kvStride);
1625+}
1626+ 
1627+template <typename T>
1628+__simd_vf__ void GetKVPhyAddrVFTndImpl(__ubuf__ uint32_t *kvPhyAddrUb, __ubuf__ int32_t *sparseIdxUb,
1629+ const uint16_t s2Loop, uint32_t s2Tail,
1630+ const uint32_t sparseBlockSize, const uint32_t kvDim,
1631+ const uint32_t kvPrefix)
1632+{
1633+ static const uint16_t s2_num_per_loop = 128;
1634+ static const uint16_t s2_num_per_reg = 64;
1635+ static const uint16_t out_offset_per_loop = 256;
1636+ static const uint16_t out_offset_per_reg = 128;
1637+ static const uint32_t invalid_value = 0xFFFFFFFF;
1638+ MicroAPI::MaskReg preg_all_b32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
1639+ MicroAPI::MaskReg preg_tail_neg_1_b32;
1640+ MicroAPI::MaskReg preg_tail_neg_2_b32;
1641+ 
1642+ MicroAPI::RegTensor<uint32_t> vreg_sparse_idx_1;
1643+ MicroAPI::RegTensor<uint32_t> vreg_sparse_idx_2;
1644+ MicroAPI::RegTensor<uint32_t> vreg_kv_prefix;
1645+ MicroAPI::RegTensor<uint32_t> vreg_kv_dim;
1646+ MicroAPI::RegTensor<uint32_t> vreg_sum_1;
1647+ MicroAPI::RegTensor<uint32_t> vreg_sum_2;
1648+ MicroAPI::RegTensor<uint32_t> vreg_mul_overflow_l_1;
1649+ MicroAPI::RegTensor<uint32_t> vreg_mul_overflow_l_2;
1650+ MicroAPI::RegTensor<uint32_t> vreg_total_offset_l_1;
1651+ MicroAPI::RegTensor<uint32_t> vreg_total_offset_h_1;
1652+ MicroAPI::RegTensor<uint32_t> vreg_total_offset_l_2;
1653+ MicroAPI::RegTensor<uint32_t> vreg_total_offset_h_2;
1654+ 
1655+ MicroAPI::Duplicate(vreg_kv_prefix, kvPrefix);
1656+ MicroAPI::Duplicate(vreg_kv_dim, kvDim);
1657+ 
1658+ for (; s2Loop > 1;) {
1659+ for (uint16_t i = 0; i < s2Loop - 1; i++) {
1660+ MicroAPI::LoadAlign<int32_t, MicroAPI::LoadDist::DIST_NORM>(
1661+ (MicroAPI::RegTensor<int32_t> &)vreg_sparse_idx_1, sparseIdxUb + i * s2_num_per_loop);
1662+ MicroAPI::LoadAlign<int32_t, MicroAPI::LoadDist::DIST_NORM>(
1663+ (MicroAPI::RegTensor<int32_t> &)vreg_sparse_idx_2, sparseIdxUb + s2_num_per_reg + i * s2_num_per_loop);
1664+ // * sparseBlockSize
1665+ MicroAPI::Muls(vreg_sparse_idx_1, vreg_sparse_idx_1, sparseBlockSize, preg_all_b32);
1666+ MicroAPI::Muls(vreg_sparse_idx_2, vreg_sparse_idx_2, sparseBlockSize, preg_all_b32);
1667+ // (kvPrefix + sparseIdx) * kvDim -> int64 物理地址
1668+ MicroAPI::Add(vreg_sum_1, vreg_sparse_idx_1, vreg_kv_prefix, preg_all_b32);
1669+ MicroAPI::Add(vreg_sum_2, vreg_sparse_idx_2, vreg_kv_prefix, preg_all_b32);
1670+ // 带进位乘法
1671+ MicroAPI::Mull(vreg_total_offset_l_1, vreg_total_offset_h_1, vreg_sum_1, vreg_kv_dim, preg_all_b32);
1672+ MicroAPI::Mull(vreg_total_offset_l_2, vreg_total_offset_h_2, vreg_sum_2, vreg_kv_dim, preg_all_b32);
1673+ // 搬出
1674+ MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(
1675+ kvPhyAddrUb + i * out_offset_per_loop, vreg_total_offset_l_1, vreg_total_offset_h_1, preg_all_b32);
1676+ MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(
1677+ kvPhyAddrUb + out_offset_per_reg + i * out_offset_per_loop, vreg_total_offset_l_2,
1678+ vreg_total_offset_h_2, preg_all_b32);
1679+ }
1680+ break;
1681+ }
1682+ 
1683+ for (uint16_t i = s2Loop - 1; i < s2Loop; i++) {
1684+ MicroAPI::MaskReg preg_tail_1_b32 = MicroAPI::UpdateMask<int32_t>(s2Tail);
1685+ MicroAPI::MaskReg preg_tail_2_b32 = MicroAPI::UpdateMask<int32_t>(s2Tail);
1686+ MicroAPI::Not(preg_tail_neg_1_b32, preg_tail_1_b32, preg_all_b32);
1687+ MicroAPI::Not(preg_tail_neg_2_b32, preg_tail_2_b32, preg_all_b32);
1688+ 
1689+ MicroAPI::LoadAlign<int32_t, MicroAPI::LoadDist::DIST_NORM>((MicroAPI::RegTensor<int32_t> &)vreg_sparse_idx_1,
1690+ sparseIdxUb + i * s2_num_per_loop);
1691+ MicroAPI::LoadAlign<int32_t, MicroAPI::LoadDist::DIST_NORM>((MicroAPI::RegTensor<int32_t> &)vreg_sparse_idx_2,
1692+ sparseIdxUb + s2_num_per_reg + i * s2_num_per_loop);
1693+ // * sparseBlockSize
1694+ MicroAPI::Muls(vreg_sparse_idx_1, vreg_sparse_idx_1, sparseBlockSize, preg_tail_1_b32);
1695+ MicroAPI::Muls(vreg_sparse_idx_2, vreg_sparse_idx_2, sparseBlockSize, preg_tail_2_b32);
1696+ // (kvPrefix + sparseIdx) * kvDim -> int64 物理地址
1697+ MicroAPI::Add(vreg_sum_1, vreg_sparse_idx_1, vreg_kv_prefix, preg_tail_1_b32);
1698+ MicroAPI::Add(vreg_sum_2, vreg_sparse_idx_2, vreg_kv_prefix, preg_tail_2_b32);
1699+ // 带进位乘法
1700+ MicroAPI::Mull(vreg_total_offset_l_1, vreg_total_offset_h_1, vreg_sum_1, vreg_kv_dim, preg_tail_1_b32);
1701+ MicroAPI::Mull(vreg_total_offset_l_2, vreg_total_offset_h_2, vreg_sum_2, vreg_kv_dim, preg_tail_2_b32);
1702+ // 无效值填充-1(0xFFFFFFFF)
1703+ MicroAPI::Duplicate<uint32_t, MicroAPI::MaskMergeMode::MERGING>(vreg_total_offset_l_1, invalid_value,
1704+ preg_tail_neg_1_b32);
1705+ MicroAPI::Duplicate<uint32_t, MicroAPI::MaskMergeMode::MERGING>(vreg_total_offset_h_1, invalid_value,
1706+ preg_tail_neg_1_b32);
1707+ MicroAPI::Duplicate<uint32_t, MicroAPI::MaskMergeMode::MERGING>(vreg_total_offset_l_2, invalid_value,
1708+ preg_tail_neg_2_b32);
1709+ MicroAPI::Duplicate<uint32_t, MicroAPI::MaskMergeMode::MERGING>(vreg_total_offset_h_2, invalid_value,
1710+ preg_tail_neg_2_b32);
1711+ MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(
1712+ kvPhyAddrUb + i * out_offset_per_loop, vreg_total_offset_l_1, vreg_total_offset_h_1, preg_all_b32);
1713+ MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(
1714+ kvPhyAddrUb + out_offset_per_reg + i * out_offset_per_loop, vreg_total_offset_l_2,
1715+ vreg_total_offset_h_2, preg_all_b32);
1716+ }
1717+}
1718+ 
1719+template <typename T>
1720+__aicore__ inline void GetKVPhyAddrVFTnd(LocalTensor<uint32_t> kvPhyAddrTensor, LocalTensor<int32_t> sparseIdxTensor,
1721+ const uint16_t s2Loop, const uint32_t s2Tail,
1722+ const uint32_t sparseBlockSize, const uint32_t kvDim,
1723+ const uint32_t kvPrefix)
1724+{
1725+ __ubuf__ uint32_t *kv_phy_addr_ub = (__ubuf__ uint32_t *)(kvPhyAddrTensor.GetPhyAddr());
1726+ __ubuf__ int32_t *sparse_idx_ub = (__ubuf__ int32_t *)(sparseIdxTensor.GetPhyAddr());
1727+ GetKVPhyAddrVFTndImpl<uint32_t>(kv_phy_addr_ub, sparse_idx_ub, s2Loop, s2Tail,
1728+ sparseBlockSize, kvDim, kvPrefix);
1729+}
1730+ 
1731+template <typename T>
1732+__simd_vf__ void GetKVPhyAddrVFBsndImpl(__ubuf__ uint32_t *kvPhyAddrUb, __ubuf__ int32_t *sparseIdxUb,
1733+ const uint16_t s2Loop, uint32_t s2Tail,
1734+ const uint32_t sparseBlockSize, const uint32_t kvDim,
1735+ const uint32_t bS2BaseLow, const uint32_t bS2BaseHigh)
1736+{
1737+ static const uint16_t s2_num_per_loop = 128;
1738+ static const uint16_t s2_num_per_reg = 64;
1739+ static const uint16_t out_offset_per_loop = 256;
1740+ static const uint16_t out_offset_per_reg = 128;
1741+ static const uint32_t invalid_value = 0xFFFFFFFF;
1742+ MicroAPI::MaskReg preg_all_b32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
1743+ MicroAPI::MaskReg add_carry_l_1;
1744+ MicroAPI::MaskReg add_carry_h_1;
1745+ MicroAPI::MaskReg add_carry_l_2;
1746+ MicroAPI::MaskReg add_carry_h_2;
1747+ MicroAPI::MaskReg preg_tail_neg_1_b32;
1748+ MicroAPI::MaskReg preg_tail_neg_2_b32;
1749+ 
1750+ MicroAPI::RegTensor<uint32_t> vreg_sparse_idx_1;
1751+ MicroAPI::RegTensor<uint32_t> vreg_sparse_idx_2;
1752+ MicroAPI::RegTensor<uint32_t> vreg_kv_dim;
1753+ MicroAPI::RegTensor<uint32_t> vreg_b_s2_base_low;
1754+ MicroAPI::RegTensor<uint32_t> vreg_b_s2_base_high;
1755+ MicroAPI::RegTensor<uint32_t> vreg_s2_offset_l_1;
1756+ MicroAPI::RegTensor<uint32_t> vreg_s2_offset_l_2;
1757+ MicroAPI::RegTensor<uint32_t> vreg_mul_overflow_l_1;
1758+ MicroAPI::RegTensor<uint32_t> vreg_mul_overflow_l_2;
1759+ MicroAPI::RegTensor<uint32_t> vreg_total_offset_l_1;
1760+ MicroAPI::RegTensor<uint32_t> vreg_total_offset_h_1;
1761+ MicroAPI::RegTensor<uint32_t> vreg_total_offset_l_2;
1762+ MicroAPI::RegTensor<uint32_t> vreg_total_offset_h_2;
1763+ MicroAPI::RegTensor<uint32_t> vreg_zero;
1764+ 
1765+ MicroAPI::Duplicate(vreg_zero, 0);
1766+ MicroAPI::Duplicate(vreg_kv_dim, kvDim);
1767+ MicroAPI::Duplicate(vreg_b_s2_base_low, bS2BaseLow);
1768+ MicroAPI::Duplicate(vreg_b_s2_base_high, bS2BaseHigh);
1769+ 
1770+ for (; s2Loop > 1;) {
1771+ for (uint16_t i = 0; i < s2Loop - 1; i++) {
1772+ MicroAPI::LoadAlign<int32_t, MicroAPI::LoadDist::DIST_NORM>(
1773+ (MicroAPI::RegTensor<int32_t> &)vreg_sparse_idx_1, sparseIdxUb + i * s2_num_per_loop);
1774+ MicroAPI::LoadAlign<int32_t, MicroAPI::LoadDist::DIST_NORM>(
1775+ (MicroAPI::RegTensor<int32_t> &)vreg_sparse_idx_2, sparseIdxUb + s2_num_per_reg + i * s2_num_per_loop);
1776+ // * sparseBlockSize
1777+ MicroAPI::Muls(vreg_sparse_idx_1, vreg_sparse_idx_1, sparseBlockSize, preg_all_b32);
1778+ MicroAPI::Muls(vreg_sparse_idx_2, vreg_sparse_idx_2, sparseBlockSize, preg_all_b32);
1779+ // sparseIdx * kvDim (带进位乘法)
1780+ MicroAPI::Mull(vreg_s2_offset_l_1, vreg_mul_overflow_l_1, vreg_sparse_idx_1, vreg_kv_dim, preg_all_b32);
1781+ MicroAPI::Mull(vreg_s2_offset_l_2, vreg_mul_overflow_l_2, vreg_sparse_idx_2, vreg_kv_dim, preg_all_b32);
1782+ // s2_offset + bS2Base (int64 + int64)
1783+ MicroAPI::Add(add_carry_l_1, vreg_total_offset_l_1, vreg_s2_offset_l_1, vreg_b_s2_base_low, preg_all_b32);
1784+ MicroAPI::Add(add_carry_l_2, vreg_total_offset_l_2, vreg_s2_offset_l_2, vreg_b_s2_base_low, preg_all_b32);
1785+ MicroAPI::AddC(add_carry_h_1, vreg_total_offset_h_1, vreg_mul_overflow_l_1, vreg_b_s2_base_high,
1786+ add_carry_l_1, preg_all_b32);
1787+ MicroAPI::AddC(add_carry_h_2, vreg_total_offset_h_2, vreg_mul_overflow_l_2, vreg_b_s2_base_high,
1788+ add_carry_l_2, preg_all_b32);
1789+ // 搬出
1790+ MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(
1791+ kvPhyAddrUb + i * out_offset_per_loop, vreg_total_offset_l_1, vreg_total_offset_h_1, preg_all_b32);
1792+ MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(
1793+ kvPhyAddrUb + out_offset_per_reg + i * out_offset_per_loop, vreg_total_offset_l_2,
1794+ vreg_total_offset_h_2, preg_all_b32);
1795+ }
1796+ break;
1797+ }
1798+ 
1799+ for (uint16_t i = s2Loop - 1; i < s2Loop; i++) {
1800+ MicroAPI::MaskReg preg_tail_1_b32 = MicroAPI::UpdateMask<int32_t>(s2Tail);
1801+ MicroAPI::MaskReg preg_tail_2_b32 = MicroAPI::UpdateMask<int32_t>(s2Tail);
1802+ MicroAPI::Not(preg_tail_neg_1_b32, preg_tail_1_b32, preg_all_b32);
1803+ MicroAPI::Not(preg_tail_neg_2_b32, preg_tail_2_b32, preg_all_b32);
1804+ 
1805+ MicroAPI::LoadAlign<int32_t, MicroAPI::LoadDist::DIST_NORM>((MicroAPI::RegTensor<int32_t> &)vreg_sparse_idx_1,
1806+ sparseIdxUb + i * s2_num_per_loop);
1807+ MicroAPI::LoadAlign<int32_t, MicroAPI::LoadDist::DIST_NORM>((MicroAPI::RegTensor<int32_t> &)vreg_sparse_idx_2,
1808+ sparseIdxUb + s2_num_per_reg + i * s2_num_per_loop);
1809+ // * sparseBlockSize
1810+ MicroAPI::Muls(vreg_sparse_idx_1, vreg_sparse_idx_1, sparseBlockSize, preg_tail_1_b32);
1811+ MicroAPI::Muls(vreg_sparse_idx_2, vreg_sparse_idx_2, sparseBlockSize, preg_tail_2_b32);
1812+ // sparseIdx * kvDim (带进位乘法)
1813+ MicroAPI::Mull(vreg_s2_offset_l_1, vreg_mul_overflow_l_1, vreg_sparse_idx_1, vreg_kv_dim, preg_tail_1_b32);
1814+ MicroAPI::Mull(vreg_s2_offset_l_2, vreg_mul_overflow_l_2, vreg_sparse_idx_2, vreg_kv_dim, preg_tail_2_b32);
1815+ // s2_offset + bS2Base (int64 + int64)
1816+ MicroAPI::Add(add_carry_l_1, vreg_total_offset_l_1, vreg_s2_offset_l_1, vreg_b_s2_base_low, preg_tail_1_b32);
1817+ MicroAPI::Add(add_carry_l_2, vreg_total_offset_l_2, vreg_s2_offset_l_2, vreg_b_s2_base_low, preg_tail_2_b32);
1818+ MicroAPI::AddC(add_carry_h_1, vreg_total_offset_h_1, vreg_mul_overflow_l_1, vreg_b_s2_base_high,
1819+ add_carry_l_1, preg_tail_1_b32);
1820+ MicroAPI::AddC(add_carry_h_2, vreg_total_offset_h_2, vreg_mul_overflow_l_2, vreg_b_s2_base_high,
1821+ add_carry_l_2, preg_tail_2_b32);
1822+ // 无效值填充-1(0xFFFFFFFF)
1823+ MicroAPI::Duplicate<uint32_t, MicroAPI::MaskMergeMode::MERGING>(vreg_total_offset_l_1, invalid_value,
1824+ preg_tail_neg_1_b32);
1825+ MicroAPI::Duplicate<uint32_t, MicroAPI::MaskMergeMode::MERGING>(vreg_total_offset_h_1, invalid_value,
1826+ preg_tail_neg_1_b32);
1827+ MicroAPI::Duplicate<uint32_t, MicroAPI::MaskMergeMode::MERGING>(vreg_total_offset_l_2, invalid_value,
1828+ preg_tail_neg_2_b32);
1829+ MicroAPI::Duplicate<uint32_t, MicroAPI::MaskMergeMode::MERGING>(vreg_total_offset_h_2, invalid_value,
1830+ preg_tail_neg_2_b32);
1831+ MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(
1832+ kvPhyAddrUb + i * out_offset_per_loop, vreg_total_offset_l_1, vreg_total_offset_h_1, preg_all_b32);
1833+ MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(
1834+ kvPhyAddrUb + out_offset_per_reg + i * out_offset_per_loop, vreg_total_offset_l_2,
1835+ vreg_total_offset_h_2, preg_all_b32);
1836+ }
1837+}
1838+ 
1839+template <typename T>
1840+__aicore__ inline void GetKVPhyAddrVFBsnd(LocalTensor<uint32_t> kvPhyAddrTensor, LocalTensor<int32_t> sparseIdxTensor,
1841+ const uint16_t s2Loop, const uint32_t s2Tail,
1842+ const uint32_t sparseBlockSize, const uint32_t kvDim,
1843+ const uint32_t bS2BaseLow, const uint32_t bS2BaseHigh)
1844+{
1845+ __ubuf__ uint32_t *kv_phy_addr_ub = (__ubuf__ uint32_t *)(kvPhyAddrTensor.GetPhyAddr());
1846+ __ubuf__ int32_t *sparse_idx_ub = (__ubuf__ int32_t *)(sparseIdxTensor.GetPhyAddr());
1847+ GetKVPhyAddrVFBsndImpl<uint32_t>(kv_phy_addr_ub, sparse_idx_ub, s2Loop, s2Tail,
1848+ sparseBlockSize, kvDim, bS2BaseLow, bS2BaseHigh);
1849+}
1850+ 
1851+TEMPLATES_DEF_NO_DEFAULT
1852+__aicore__ inline int32_t CSABlockVec<TEMPLATE_ARGS>::GetSeqLen(int32_t bIdx, bool hasActualSeq, bool hasCuSeqlens,
1853+ GlobalTensor<int32_t> &actualSeqGm,
1854+ GlobalTensor<int32_t> &cuSeqlensGm, int64_t defaultSize)
1855+{
1856+ if (hasActualSeq) {
1857+ return actualSeqGm.GetValue(bIdx);
1858+ } else if (hasCuSeqlens) {
1859+ return cuSeqlensGm.GetValue(bIdx + 1) - cuSeqlensGm.GetValue(bIdx);
1860+ } else {
1861+ return defaultSize;
1862+ }
1863+}
1864+ 
1865+TEMPLATES_DEF_NO_DEFAULT
1866+__aicore__ inline int32_t CSABlockVec<TEMPLATE_ARGS>::CalcCurValidS2(
1867+ uint32_t bIdx, int32_t s1Idx, int32_t actualS1Size, bool isOriKv, bool hasActualSeqKvlen, bool hasCuSeqlensKv,
1868+ GlobalTensor<int32_t> &cuSeqlensQGm, GlobalTensor<int32_t> &actualSeqKvlenGm,
1869+ GlobalTensor<int32_t> &cuSeqlensKvGm, GlobalTensor<int32_t> &topkLengthGm,
1870+ GlobalTensor<int32_t> &cmpResidualKvGm, ConstInfo &constInfo, int32_t sparseBlockCount)
1871+{
1872+ int32_t curValidS2 = 0;
1873+ if (isOriKv) {
1874+ if constexpr (TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||
1875+ TEMPLATE_MODE == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) {
1876+ uint64_t topkIdx = (LAYOUT_T == SMLA_LAYOUT::TND) ? (cuSeqlensQGm.GetValue(bIdx) + s1Idx) :
1877+ (bIdx * constInfo.s1Size + s1Idx);
1878+ int32_t topkLen = constInfo.hasOriTopkLength ? topkLengthGm.GetValue(topkIdx) : sparseBlockCount;
1879+ curValidS2 = Min(topkLen, sparseBlockCount);
1880+ }
1881+ } else {
1882+ if constexpr (TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE) {
1883+ if (constInfo.cmpMaskMode == 0) {
1884+ uint64_t topkIdx = (LAYOUT_T == SMLA_LAYOUT::TND) ? (cuSeqlensQGm.GetValue(bIdx) + s1Idx) :
1885+ (bIdx * constInfo.s1Size + s1Idx);
1886+ int32_t topkLen = constInfo.hasCmpTopkLength ? topkLengthGm.GetValue(topkIdx) : sparseBlockCount;
1887+ curValidS2 = Min(topkLen, sparseBlockCount);
1888+ } else {
1889+ int32_t actualCmpS2Size = GetSeqLen(bIdx, hasActualSeqKvlen, hasCuSeqlensKv, actualSeqKvlenGm,
1890+ cuSeqlensKvGm, constInfo.s2Size);
1891+ int32_t cmpRestoredSize =
1892+ actualCmpS2Size * static_cast<int32_t>(constInfo.cmpRatio) + cmpResidualKvGm.GetValue(bIdx);
1893+ int32_t numerator = cmpRestoredSize - actualS1Size + 1 + s1Idx;
1894+ curValidS2 =
1895+ (numerator > 0) ? Min(sparseBlockCount, numerator / static_cast<int32_t>(constInfo.cmpRatio)) : 0;
1896+ }
1897+ } else if constexpr (TEMPLATE_MODE == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) {
1898+ uint64_t topkIdx = (LAYOUT_T == SMLA_LAYOUT::TND) ? (cuSeqlensQGm.GetValue(bIdx) + s1Idx) :
1899+ (bIdx * constInfo.s1Size + s1Idx);
1900+ int32_t topkLen = constInfo.hasCmpTopkLength ? topkLengthGm.GetValue(topkIdx) : sparseBlockCount;
1901+ curValidS2 = Min(topkLen, sparseBlockCount);
1902+ }
1903+ }
1904+ return curValidS2;
1905+}
1906+ 
1907+TEMPLATES_DEF_NO_DEFAULT
1908+__aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::CopyPhyAddrToGm(LocalTensor<uint32_t> kvPhyAddrUb, int64_t bS1Idx,
1909+ int64_t s1Idx, int64_t validS2, int64_t alignNum,
1910+ GlobalTensor<uint32_t> &phyAddrGm,
1911+ uint32_t alignedSparseBlockCount)
1912+{
1913+ constexpr int64_t numPerBlock = 32;
1914+ DataCopyParams dataCopyParams;
1915+ dataCopyParams.blockCount = 1U;
1916+ dataCopyParams.blockLen = ((validS2 + alignNum - 1) / alignNum * alignNum) * sizeof(int64_t) / numPerBlock;
1917+ dataCopyParams.srcGap = 0U;
1918+ dataCopyParams.dstGap = 0U;
1919+ DataCopy(phyAddrGm[(bS1Idx + s1Idx) * alignedSparseBlockCount * 2], kvPhyAddrUb, dataCopyParams);
1920+}
1921+ 
1922+TEMPLATES_DEF_NO_DEFAULT
1923+__aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::CopyPaTableToUb(LocalTensor<int32_t> blkTableUb, int64_t bIdx,
1924+ GlobalTensor<int32_t> &blockTableGm,
1925+ uint32_t maxBlockNumPerBatch)
1926+{
1927+ DataCopyExtParams dataCopyParams;
1928+ dataCopyParams.blockCount = 1U;
1929+ dataCopyParams.blockLen = maxBlockNumPerBatch * sizeof(int32_t);
1930+ dataCopyParams.srcStride = 0U;
1931+ dataCopyParams.dstStride = 0U;
1932+ DataCopyPadExtParams<int32_t> padParams;
1933+ DataCopyPad(blkTableUb, blockTableGm[bIdx * maxBlockNumPerBatch], dataCopyParams, padParams);
1934+}
1935+ 
1936+TEMPLATES_DEF_NO_DEFAULT
1937+__aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::CopySparseIdxToUb(LocalTensor<int32_t> sparseIdxUb, int64_t bS1Idx,
1938+ int64_t s1Idx, int64_t validS2,
1939+ GlobalTensor<int32_t> &sparseIndicesGm,
1940+ uint32_t sparseBlockCount)
1941+{
1942+ DataCopyExtParams dataCopyParams;
1943+ dataCopyParams.blockCount = 1U;
1944+ dataCopyParams.blockLen = validS2 * sizeof(int32_t);
1945+ dataCopyParams.srcStride = 0U;
1946+ dataCopyParams.dstStride = 0U;
1947+ DataCopyPadExtParams<int32_t> padParams;
1948+ DataCopyPad(sparseIdxUb, sparseIndicesGm[(bS1Idx + s1Idx) * sparseBlockCount], dataCopyParams, padParams);
1949+}
1950+ 
1951+TEMPLATES_DEF_NO_DEFAULT
1952+__aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::GetKVPhyAddrForKvType(
1953+ uint32_t bN2StartIdx, uint32_t bN2EndIdx, uint32_t gS1StartIdx, uint32_t nextGs1Idx, bool hasActualSeqQlen,
1954+ bool hasCuSeqlensQ, bool hasActualSeqKvlen, bool hasCuSeqlensKv, GlobalTensor<int32_t> actualSeqQlenGm,
1955+ GlobalTensor<int32_t> cuSeqlensQGm, GlobalTensor<int32_t> actualSeqKvlenGm, GlobalTensor<int32_t> cuSeqlensKvGm,
1956+ GlobalTensor<int32_t> topkLengthGm, GlobalTensor<int32_t> cmpResidualKvGm, ConstInfo &constInfo,
1957+ GlobalTensor<int32_t> &blockTableGm, GlobalTensor<int32_t> &sparseIndicesGm, GlobalTensor<uint32_t> &phyAddrGm,
1958+ uint32_t kvStride, uint32_t blockSize, uint32_t maxBlockNumPerBatch, uint32_t sparseBlockCount,
1959+ uint32_t alignedSparseBlockCount, bool isOriKv)
1960+{
1961+ static constexpr uint16_t s2NumPerLoop = 128;
1962+ static constexpr uint32_t vecCoreNum = IS_SPLIT_G ? 4 : 2;
1963+ uint32_t vecCoreIdx = IS_SPLIT_G ? constInfo.aivIdx % 4 : constInfo.aivIdx % 2;
1964+ TBuf<> blkTableBuf;
1965+ TBuf<> sparseIdxBuf;
1966+ TBuf<> kvPhyAddrBuf;
1967+ int16_t shiftRightNum = 0;
1968+ LocalTensor<int32_t> blkTableUb;
1969+ 
1970+ if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND) {
1971+ int32_t blkSize = static_cast<int32_t>(blockSize);
1972+ while (blkSize > 1) {
1973+ blkSize >>= 1;
1974+ shiftRightNum++;
1975+ }
1976+ tPipe->InitBuffer(blkTableBuf, maxBlockNumPerBatch * sizeof(int32_t));
1977+ blkTableUb = blkTableBuf.template Get<int32_t>();
1978+ }
1979+ tPipe->InitBuffer(sparseIdxBuf, alignedSparseBlockCount * sizeof(int32_t));
1980+ tPipe->InitBuffer(kvPhyAddrBuf, alignedSparseBlockCount * sizeof(int64_t));
1981+ LocalTensor<int32_t> sparseIdxUb = sparseIdxBuf.template Get<int32_t>();
1982+ LocalTensor<uint32_t> kvPhyAddrUb = kvPhyAddrBuf.template Get<uint32_t>();
1983+ 
1984+ // 第一遍: 统计totalValidS1
1985+ int64_t totalValidS1 = 0;
1986+ uint32_t tmpGS1Start = gS1StartIdx;
1987+ for (uint32_t bIdx = bN2StartIdx; bIdx < bN2EndIdx; ++bIdx) {
1988+ bool lastBN = (bIdx == bN2EndIdx - 1);
1989+ int32_t actualS1Size =
1990+ GetSeqLen(bIdx, hasActualSeqQlen, hasCuSeqlensQ, actualSeqQlenGm, cuSeqlensQGm, constInfo.s1Size);
1991+ int32_t s1End = actualS1Size;
1992+ if (lastBN && nextGs1Idx != 0) {
1993+ s1End = nextGs1Idx;
1994+ }
1995+ 
1996+ int64_t bS1IdxBase = 0;
1997+ if constexpr (LAYOUT_T == SMLA_LAYOUT::TND) {
1998+ bS1IdxBase = hasCuSeqlensQ ? cuSeqlensQGm.GetValue(bIdx) : constInfo.s1Size * bIdx;
1999+ } else {
2000+ bS1IdxBase = constInfo.s1Size * bIdx;
2001+ }
2002+ 
2003+ for (int32_t s1Idx = tmpGS1Start; s1Idx < s1End; ++s1Idx) {
2004+ int32_t curValidS2 = CalcCurValidS2(bIdx, s1Idx, actualS1Size, isOriKv, hasActualSeqKvlen,
2005+ hasCuSeqlensKv, cuSeqlensQGm, actualSeqKvlenGm, cuSeqlensKvGm,
2006+ topkLengthGm, cmpResidualKvGm, constInfo,
2007+ static_cast<int32_t>(sparseBlockCount));
2008+ if (curValidS2 > 0) {
2009+ totalValidS1++;
2010+ }
2011+ }
2012+ tmpGS1Start = 0;
2013+ }
2014+ 
2015+ int64_t s1PerVecCore = totalValidS1 / vecCoreNum;
2016+ int64_t s1Tail = totalValidS1 % vecCoreNum;
2017+ int64_t curStart = s1PerVecCore * vecCoreIdx + Min((int64_t)vecCoreIdx, s1Tail);
2018+ int64_t curCount = s1PerVecCore + (vecCoreIdx < (uint32_t)s1Tail ? 1 : 0);
2019+ 
2020+ if (curCount == 0) {
2021+ return;
2022+ }
2023+ 
2024+ // 第二遍: 实际计算
2025+ int64_t validCounter = 0;
2026+ int64_t processedCount = 0;
2027+ tmpGS1Start = gS1StartIdx;
2028+ bool done = false;
2029+ 
2030+ if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND) {
2031+ SetFlag<AscendC::HardEvent::V_MTE2>(blkTableVToMte2FlagId);
2032+ }
2033+ SetFlag<AscendC::HardEvent::V_MTE2>(sparseIdxVToMte2FlagId);
2034+ SetFlag<AscendC::HardEvent::MTE3_V>(phyAddrMte3ToVFlagId);
2035+ for (uint32_t bIdx = bN2StartIdx; bIdx < bN2EndIdx && !done; ++bIdx) {
2036+ bool lastBN = (bIdx == bN2EndIdx - 1);
2037+ int32_t actualS1Size =
2038+ GetSeqLen(bIdx, hasActualSeqQlen, hasCuSeqlensQ, actualSeqQlenGm, cuSeqlensQGm, constInfo.s1Size);
2039+ int64_t bS1Idx = 0;
2040+ if constexpr (LAYOUT_T == SMLA_LAYOUT::TND) {
2041+ bS1Idx = hasCuSeqlensQ ? cuSeqlensQGm.GetValue(bIdx) : constInfo.s1Size * bIdx;
2042+ } else {
2043+ bS1Idx = constInfo.s1Size * bIdx;
2044+ }
2045+ 
2046+ int32_t s1End = actualS1Size;
2047+ if (lastBN && nextGs1Idx != 0) {
2048+ s1End = nextGs1Idx;
2049+ }
2050+ 
2051+ // per-batch 参数预计算
2052+ uint32_t kvPrefix = 0;
2053+ uint32_t bS2BaseLow = 0;
2054+ uint32_t bS2BaseHigh = 0;
2055+ if constexpr (LAYOUT_T == SMLA_LAYOUT::TND) {
2056+ kvPrefix = static_cast<uint32_t>(cuSeqlensKvGm.GetValue(bIdx));
2057+ } else {
2058+ uint32_t s2Size = isOriKv ? static_cast<uint32_t>(constInfo.s2Size) :
2059+ static_cast<uint32_t>(constInfo.cmpS2Size);
2060+ uint64_t bS2Base = static_cast<uint64_t>(bIdx) * s2Size *
2061+ static_cast<uint64_t>(constInfo.dSize);
2062+ bS2BaseLow = static_cast<uint32_t>(bS2Base);
2063+ bS2BaseHigh = static_cast<uint32_t>(bS2Base >> 32U);
2064+ }
2065+ 
2066+ if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND) {
2067+ WaitFlag<AscendC::HardEvent::V_MTE2>(blkTableVToMte2FlagId);
2068+ CopyPaTableToUb(blkTableUb, bIdx, blockTableGm, maxBlockNumPerBatch);
2069+ SetFlag<AscendC::HardEvent::MTE2_V>(blkTableMte2ToVFlagId);
2070+ WaitFlag<AscendC::HardEvent::MTE2_V>(blkTableMte2ToVFlagId);
2071+ }
2072+ 
2073+ for (int32_t s1Idx = tmpGS1Start; s1Idx < s1End; ++s1Idx) {
2074+ int32_t curValidS2 = CalcCurValidS2(bIdx, s1Idx, actualS1Size, isOriKv, hasActualSeqKvlen,
2075+ hasCuSeqlensKv, cuSeqlensQGm, actualSeqKvlenGm, cuSeqlensKvGm,
2076+ topkLengthGm, cmpResidualKvGm, constInfo,
2077+ static_cast<int32_t>(sparseBlockCount));
2078+ if (curValidS2 <= 0) {
2079+ continue;
2080+ }
2081+ 
2082+ if (validCounter < curStart || validCounter >= curStart + curCount) {
2083+ validCounter++;
2084+ continue;
2085+ }
2086+ validCounter++;
2087+ 
2088+ uint16_t s2Loop = (curValidS2 + s2NumPerLoop - 1) / s2NumPerLoop;
2089+ int32_t s2Tail = curValidS2 - (s2Loop - 1) * s2NumPerLoop;
2090+ WaitFlag<AscendC::HardEvent::V_MTE2>(sparseIdxVToMte2FlagId);
2091+ CopySparseIdxToUb(sparseIdxUb, bS1Idx, s1Idx, curValidS2, sparseIndicesGm, sparseBlockCount);
2092+ SetFlag<AscendC::HardEvent::MTE2_V>(sparseIdxMte2ToVFlagId);
2093+ 
2094+ WaitFlag<AscendC::HardEvent::MTE2_V>(sparseIdxMte2ToVFlagId);
2095+ WaitFlag<AscendC::HardEvent::MTE3_V>(phyAddrMte3ToVFlagId);
2096+ if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND) {
2097+ GetKVPhyAddrVFPa<uint32_t>(kvPhyAddrUb, sparseIdxUb, blkTableUb, s2Loop, s2Tail, blockSize,
2098+ shiftRightNum, constInfo.sparseBlockSize,
2099+ static_cast<uint32_t>(constInfo.dSize), kvStride);
2100+ } else if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::TND) {
2101+ GetKVPhyAddrVFTnd<uint32_t>(kvPhyAddrUb, sparseIdxUb, s2Loop, s2Tail,
2102+ constInfo.sparseBlockSize,
2103+ static_cast<uint32_t>(constInfo.dSize), kvPrefix);
2104+ } else {
2105+ GetKVPhyAddrVFBsnd<uint32_t>(kvPhyAddrUb, sparseIdxUb, s2Loop, s2Tail,
2106+ constInfo.sparseBlockSize,
2107+ static_cast<uint32_t>(constInfo.dSize),
2108+ bS2BaseLow, bS2BaseHigh);
2109+ }
2110+ SetFlag<AscendC::HardEvent::V_MTE2>(sparseIdxVToMte2FlagId);
2111+ SetFlag<AscendC::HardEvent::V_MTE3>(phyAddrVToMte3FlagId);
2112+ WaitFlag<AscendC::HardEvent::V_MTE3>(phyAddrVToMte3FlagId);
2113+ CopyPhyAddrToGm(kvPhyAddrUb, bS1Idx, s1Idx, curValidS2, s2NumPerLoop, phyAddrGm,
2114+ alignedSparseBlockCount);
2115+ SetFlag<AscendC::HardEvent::MTE3_V>(phyAddrMte3ToVFlagId);
2116+ 
2117+ processedCount++;
2118+ if (processedCount >= curCount) {
2119+ done = true;
2120+ break;
2121+ }
2122+ }
2123+ if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND) {
2124+ SetFlag<AscendC::HardEvent::V_MTE2>(blkTableVToMte2FlagId);
2125+ }
2126+ tmpGS1Start = 0;
2127+ }
2128+ if constexpr (KV_LAYOUT_T == SMLA_LAYOUT::PA_BBND) {
2129+ WaitFlag<AscendC::HardEvent::V_MTE2>(blkTableVToMte2FlagId);
2130+ }
2131+ WaitFlag<AscendC::HardEvent::V_MTE2>(sparseIdxVToMte2FlagId);
2132+ WaitFlag<AscendC::HardEvent::MTE3_V>(phyAddrMte3ToVFlagId);
2133+}
2134+ 
2135+TEMPLATES_DEF_NO_DEFAULT
2136+__aicore__ inline void CSABlockVec<TEMPLATE_ARGS>::GetKVPhyAddr(
2137+ uint32_t hasLoad, uint32_t bN2StartIdx, uint32_t bN2EndIdx, uint32_t gS1StartIdx, uint32_t nextGs1Idx,
2138+ bool hasActualSeqQlen, bool hasCuSeqlensQ, bool hasActualSeqOriKvlen, bool hasCuSeqlensOriKv,
2139+ GlobalTensor<int32_t> actualSeqOriKvlenGm, GlobalTensor<int32_t> cuSeqlensOriKvGm,
2140+ GlobalTensor<int32_t> oriTopkLengthGm, bool hasActualSeqCmpKvlen, bool hasCuSeqlensCmpKv,
2141+ GlobalTensor<int32_t> actualSeqCmpKvlenGm, GlobalTensor<int32_t> cuSeqlensCmpKvGm,
2142+ GlobalTensor<int32_t> cmpTopkLengthGm, GlobalTensor<int32_t> cmpResidualKvGm, GlobalTensor<int32_t> actualSeqQlenGm,
2143+ GlobalTensor<int32_t> cuSeqlensQGm, __gm__ uint8_t *workspace, ConstInfo &constInfo)
2144+{
2145+ if (hasLoad == 0) {
2146+ SyncAll();
2147+ tPipe->Reset();
2148+ return;
2149+ }
2150+ 
2151+ // GM分配: ori在前, cmp在后
2152+ int64_t v0TotalOffset = 0;
2153+ uint32_t v0ResSize = constInfo.s2BaseSize * constInfo.dSize * sizeof(Q_T);
2154+ if constexpr (IS_SPLIT_G) {
2155+ v0TotalOffset = v0ResSize * 3 * (GetBlockNum() >> 1U);
2156+ } else {
2157+ v0TotalOffset = v0ResSize * 3 * GetBlockNum();
2158+ }
2159+ 
2160+ // SMLA特有: 加上s2RealBuf大小
2161+ constexpr uint32_t TRIPLE_BUFFER_NUM = 3;
2162+ constexpr uint32_t S2_REAL_BUF_LEN = 128;
2163+ v0TotalOffset += TRIPLE_BUFFER_NUM * S2_REAL_BUF_LEN * sizeof(int32_t) * GetBlockNum();
2164+ 
2165+ uint32_t totalBS1 = (LAYOUT_T == SMLA_LAYOUT::TND) ? constInfo.s1Size : (constInfo.bSize * constInfo.s1Size);
2166+ 
2167+ uint64_t oriPhyAddrSize = 0;
2168+ if constexpr (TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||
2169+ TEMPLATE_MODE == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) {
2170+ oriPhyAddrSize = static_cast<uint64_t>(totalBS1) * constInfo.alignedOriSparseBlockCount * sizeof(int64_t);
2171+ this->oriKvPhyAddrGm.SetGlobalBuffer((__gm__ uint32_t *)(workspace + v0TotalOffset));
2172+ }
2173+ 
2174+ if constexpr (TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE ||
2175+ TEMPLATE_MODE == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) {
2176+ uint64_t cmpPhyAddrSize =
2177+ static_cast<uint64_t>(totalBS1) * constInfo.alignedCmpSparseBlockCount * sizeof(int64_t);
2178+ this->cmpKvPhyAddrGm.SetGlobalBuffer((__gm__ uint32_t *)(workspace + v0TotalOffset + oriPhyAddrSize));
2179+ }
2180+ 
2181+ // ori部分 (先计算)
2182+ if constexpr (TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||
2183+ TEMPLATE_MODE == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) {
2184+ GetKVPhyAddrForKvType(bN2StartIdx, bN2EndIdx, gS1StartIdx, nextGs1Idx, hasActualSeqQlen, hasCuSeqlensQ,
2185+ hasActualSeqOriKvlen, hasCuSeqlensOriKv, actualSeqQlenGm, cuSeqlensQGm,
2186+ actualSeqOriKvlenGm, cuSeqlensOriKvGm, oriTopkLengthGm, cmpResidualKvGm, constInfo,
2187+ oriBlockTableGm, oriSparseIndicesGm, oriKvPhyAddrGm, constInfo.oriKeyStride0,
2188+ constInfo.oriBlockSize, constInfo.oriMaxBlockNumPerBatch, constInfo.oriSparseBlockCount,
2189+ constInfo.alignedOriSparseBlockCount, true);
2190+ }
2191+ 
2192+ // cmp部分 (后计算)
2193+ if constexpr (TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE ||
2194+ TEMPLATE_MODE == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) {
2195+ GetKVPhyAddrForKvType(bN2StartIdx, bN2EndIdx, gS1StartIdx, nextGs1Idx, hasActualSeqQlen, hasCuSeqlensQ,
2196+ hasActualSeqCmpKvlen, hasCuSeqlensCmpKv, actualSeqQlenGm, cuSeqlensQGm,
2197+ actualSeqCmpKvlenGm, cuSeqlensCmpKvGm, cmpTopkLengthGm, cmpResidualKvGm, constInfo,
2198+ cmpBlockTableGm, cmpSparseIndicesGm, cmpKvPhyAddrGm, constInfo.cmpKeyStride0,
2199+ constInfo.cmpBlockSize, constInfo.cmpMaxBlockNumPerBatch, constInfo.cmpSparseBlockCount,
2200+ constInfo.alignedCmpSparseBlockCount, false);
2201+ }
2202+ 
2203+ SyncAll();
2204+ tPipe->Reset();
2205+}
2206+ 
1376TEMPLATES_DEF2207TEMPLATES_DEF
1377class CSABlockVecDummy {2208class CSABlockVecDummy {
1378public:2209public:
@@ -1381,12 +2212,12 @@ public:
1381 {2212 {
1382 }2213 }
1383 __aicore__ inline void InitGlobalBuffer(__gm__ uint8_t *oriKV, __gm__ uint8_t *cmpKV,2214 __aicore__ inline void InitGlobalBuffer(__gm__ uint8_t *oriKV, __gm__ uint8_t *cmpKV,
1384- __gm__ uint8_t *oriSparseIndices, __gm__ uint8_t *cmpSparseIndices, __gm__ uint8_t *oriBlockTable,2215+ __gm__ uint8_t *oriSparseIndices, __gm__ uint8_t *cmpSparseIndices, __gm__ uint8_t *oriBlockTable,
1385- __gm__ uint8_t *cmpBlockTable, __gm__ uint8_t *sequsedQ, __gm__ uint8_t *sinks,2216+ __gm__ uint8_t *cmpBlockTable, __gm__ uint8_t *sequsedQ, __gm__ uint8_t *sinks,
1386- __gm__ uint8_t *sequsedOriKv, __gm__ uint8_t *sequsedCmpKv, __gm__ uint8_t *cmpResidualKv) {}2217+ __gm__ uint8_t *sequsedOriKv, __gm__ uint8_t *sequsedCmpKv, __gm__ uint8_t *cmpResidualKv) {}
1387 __aicore__ inline void InitVecBlock(TPipe *pipe, __gm__ uint8_t *cuSeqlensQ,2218 __aicore__ inline void InitVecBlock(TPipe *pipe, __gm__ uint8_t *cuSeqlensQ,
1388- __gm__ uint8_t *cuSeqlensOriKv, __gm__ uint8_t *cuSeqlensCmpKv, __gm__ uint8_t *seqUsedOriKV,2219+ __gm__ uint8_t *cuSeqlensOriKv, __gm__ uint8_t *cuSeqlensCmpKv, __gm__ uint8_t *seqUsedOriKV,
1389- __gm__ uint8_t *seqUsedCmpKV, __gm__ uint8_t *cmpResidualKV) {};2220+ __gm__ uint8_t *seqUsedCmpKV, __gm__ uint8_t *cmpResidualKV) {};
1390 __aicore__ inline void InitS2SplitStaging(2221 __aicore__ inline void InitS2SplitStaging(
1391 Buffer<BufferType::GM, SyncType::INNER_CORE_SYNC> &fdStaging) {}2222 Buffer<BufferType::GM, SyncType::INNER_CORE_SYNC> &fdStaging) {}
1392 __aicore__ inline void InitS2SplitStaging(2223 __aicore__ inline void InitS2SplitStaging(
Mattention/sparse_flash_mla/op_kernel/arch35/sparse_flash_mla_csa_kernel.h+107-60
@@ -37,6 +37,9 @@
37#else37#else
38#include "../common/CopyInL1.h"38#include "../common/CopyInL1.h"
39#endif39#endif
40+#if __has_include("common/buffers_policy_3buff_sfa.h")
41+#include "common/buffers_policy_3buff_sfa.h"
42+#endif
40 43 
41#include "kernel_operator_list_tensor_intf.h"44#include "kernel_operator_list_tensor_intf.h"
42 45 
@@ -104,7 +107,9 @@ private:
104 const SparseFlashMlaTilingData *__restrict tilingData;107 const SparseFlashMlaTilingData *__restrict tilingData;
105 /* 编译期常量的基本块信息 */108 /* 编译期常量的基本块信息 */
106 static constexpr uint64_t SYNC_MODE = 4;109 static constexpr uint64_t SYNC_MODE = 4;
107- static constexpr uint32_t PRELOAD_NUM = 2;110+ static constexpr uint32_t PRELOAD_NUM = 3;
111+ static constexpr uint32_t crossCoreMte2SyncFlagId = 15; // IS_SPLIT_G 核间 MTE2 同步 flag ID
112+ static constexpr uint32_t SPARSE_BLOCK_ALIGN_NUM = 128;
108 113 
109 uint32_t crossCoreSyncBufId = 0;114 uint32_t crossCoreSyncBufId = 0;
110 /* 核间通道 */115 /* 核间通道 */
@@ -117,7 +122,6 @@ private:
117 // mm2左矩阵P122 // mm2左矩阵P
118 BufferManager<BufferType::L1> l1BufferManager;123 BufferManager<BufferType::L1> l1BufferManager;
119 BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> l1PBuffers;124 BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> l1PBuffers;
120- BuffersPolicy3buff<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> l1RightBuffers;
121 GlobalTensor<uint32_t> metadataGm;125 GlobalTensor<uint32_t> metadataGm;
122 GlobalTensor<int32_t> cuSeqlensQGm;126 GlobalTensor<int32_t> cuSeqlensQGm;
123 GlobalTensor<int32_t> cuSeqlensOriKvGm;127 GlobalTensor<int32_t> cuSeqlensOriKvGm;
@@ -136,7 +140,7 @@ private:
136 bool hasActualSeqOriKvlen = false;140 bool hasActualSeqOriKvlen = false;
137 bool hasActualSeqCmpKvlen = false;141 bool hasActualSeqCmpKvlen = false;
138 /* workspace 空间 */142 /* workspace 空间 */
139- BuffersPolicy3buff<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> v0ResGmBuffers;143+ BuffersPolicy3buffSFA<BufferType::GM, SyncType::CROSS_CORE_SYNC_BACKWARD> v0ResGmBuffers;
140 BufferManager<BufferType::GM> fdStagingBufferManager;144 BufferManager<BufferType::GM> fdStagingBufferManager;
141 BuffersPolicySingleBuffer<BufferType::GM> fdStagingBuffer;145 BuffersPolicySingleBuffer<BufferType::GM> fdStagingBuffer;
142 BuffersPolicySingleBuffer<BufferType::GM> intraCoreCombineBuffer;146 BuffersPolicySingleBuffer<BufferType::GM> intraCoreCombineBuffer;
@@ -144,6 +148,13 @@ private:
144 /* 核Index信息 */148 /* 核Index信息 */
145 int32_t aicIdx;149 int32_t aicIdx;
146 150 
151+ /* Init阶段metadata解析结果 */
152+ uint32_t bN2StartIdx;
153+ uint32_t gS1StartIdx;
154+ uint32_t bN2EndIdx;
155+ uint32_t nextGs1Idx;
156+ uint32_t hasLoad;
157+ 
147 /* 初始化后不变的信息 */158 /* 初始化后不变的信息 */
148 ConstInfo constInfo;159 ConstInfo constInfo;
149 160 
@@ -187,18 +198,42 @@ __aicore__ inline void SparseFlashMlaCsaKernel<CubeBlockType, VecBlockType>::Ini
187 cmpResidualKV);198 cmpResidualKV);
188 vecBlock.InitVecBlock(tPipe, cuSeqlensQ, cuSeqlensOriKv, cuSeqlensCmpKv, seqUsedOriKV, seqUsedCmpKV, cmpResidualKV);199 vecBlock.InitVecBlock(tPipe, cuSeqlensQ, cuSeqlensOriKv, cuSeqlensCmpKv, seqUsedOriKV, seqUsedCmpKV, cmpResidualKV);
189 vecBlock.CleanOutput(attentionOut, softmaxLse, constInfo);200 vecBlock.CleanOutput(attentionOut, softmaxLse, constInfo);
201+ 
202+ // 从meta data解析分核信息
203+ bN2StartIdx = metadataGm.GetValue(GetAttrAbsIndex(aicIdx, FA_BN2_START_INDEX, false));
204+ gS1StartIdx = metadataGm.GetValue(GetAttrAbsIndex(aicIdx, FA_M_START_INDEX, false));
205+ bN2EndIdx = metadataGm.GetValue(GetAttrAbsIndex(aicIdx, FA_BN2_END_INDEX, false));
206+ nextGs1Idx = metadataGm.GetValue(GetAttrAbsIndex(aicIdx, FA_M_END_INDEX, false));
207+ hasLoad = metadataGm.GetValue(GetAttrAbsIndex(aicIdx, FA_CORE_ENABLE_INDEX, false));
208+ if (nextGs1Idx != 0) {
209+ bN2EndIdx++;
210+ }
211+ 
212+ this->InitGlobalBuffer(query, oriKV, cmpKV, oriSparseIndices, cmpSparseIndices, oriBlockTable, cmpBlockTable,
213+ cuSeqlensQ, cuSeqlensOriKv, cuSeqlensCmpKv, sequsedQ, seqUsedOriKV, seqUsedCmpKV,
214+ cmpResidualKV, oriTopkLength, cmpTopkLength, sinks, workspace, tiling, tPipe); // gm设置
215+ 
216+ if ASCEND_IS_AIV {
217+ if constexpr ((TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE ||
218+ TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||
219+ TEMPLATE_MODE == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) &&
220+ IS_VEC_S2PHYADDR) {
221+ this->vecBlock.GetKVPhyAddr(hasLoad, bN2StartIdx, bN2EndIdx, gS1StartIdx, nextGs1Idx, hasActualSeqQlen,
222+ hasCuSeqlensQ, hasActualSeqOriKvlen, hasCuSeqlensOriKv, actualSeqOriKvlenGm,
223+ cuSeqlensOriKvGm, oriTopkLengthGm, hasActualSeqCmpKvlen, hasCuSeqlensCmpKv,
224+ actualSeqCmpKvlenGm, cuSeqlensCmpKvGm, cmpTopkLengthGm, cmpResidualKvGm,
225+ actualSeqQlenGm, cuSeqlensQGm, workspace, constInfo);
226+ }
227+ }
228+ 
190 InitMMResBuf(workspace);229 InitMMResBuf(workspace);
191 if constexpr (IS_BATCH_CONSISTENCY) {230 if constexpr (IS_BATCH_CONSISTENCY) {
192 vecBlock.InitS2SplitStaging(intraCoreCombineBuffer.Get(), crossCoreCombineBuffer.Get());231 vecBlock.InitS2SplitStaging(intraCoreCombineBuffer.Get(), crossCoreCombineBuffer.Get());
193 } else {232 } else {
194 vecBlock.InitS2SplitStaging(fdStagingBuffer.Get());233 vecBlock.InitS2SplitStaging(fdStagingBuffer.Get());
195 }234 }
196- cubeBlock.InitCubeBlock(pipe, l1BufferManager, query);235+ cubeBlock.InitCubeBlock(pipe, query);
197 this->ComputeConstexpr();236 this->ComputeConstexpr();
198- 
199- this->InitGlobalBuffer(query, oriKV, cmpKV, oriSparseIndices, cmpSparseIndices, oriBlockTable, cmpBlockTable,
200- cuSeqlensQ, cuSeqlensOriKv, cuSeqlensCmpKv, sequsedQ, seqUsedOriKV, seqUsedCmpKV,
201- cmpResidualKV, oriTopkLength, cmpTopkLength, sinks, workspace, tiling, tPipe); // gm设置
202 this->InitLocalBuffer();237 this->InitLocalBuffer();
203}238}
204 239 
@@ -232,6 +267,10 @@ __aicore__ inline void SparseFlashMlaCsaKernel<CubeBlockType, VecBlockType>::Par
232 constInfo.cmpS2Size = sparseFlashMLACmpParams.cmpKvSeqSize;267 constInfo.cmpS2Size = sparseFlashMLACmpParams.cmpKvSeqSize;
233 constInfo.oriSparseBlockCount = sparseFlashMLABaseParams.oriSparseBlockCount;268 constInfo.oriSparseBlockCount = sparseFlashMLABaseParams.oriSparseBlockCount;
234 constInfo.cmpSparseBlockCount = sparseFlashMLACmpParams.cmpSparseBlockCount;269 constInfo.cmpSparseBlockCount = sparseFlashMLACmpParams.cmpSparseBlockCount;
270+ constInfo.alignedOriSparseBlockCount =
271+ (constInfo.oriSparseBlockCount + SPARSE_BLOCK_ALIGN_NUM - 1) / SPARSE_BLOCK_ALIGN_NUM * SPARSE_BLOCK_ALIGN_NUM;
272+ constInfo.alignedCmpSparseBlockCount =
273+ (constInfo.cmpSparseBlockCount + SPARSE_BLOCK_ALIGN_NUM - 1) / SPARSE_BLOCK_ALIGN_NUM * SPARSE_BLOCK_ALIGN_NUM;
235 constInfo.cmpRatio = sparseFlashMLACmpParams.cmpRatio;274 constInfo.cmpRatio = sparseFlashMLACmpParams.cmpRatio;
236 constInfo.oriMaskMode = sparseFlashMLABaseParams.oriMaskMode;275 constInfo.oriMaskMode = sparseFlashMLABaseParams.oriMaskMode;
237 constInfo.cmpMaskMode = sparseFlashMLACmpParams.cmpMaskMode;276 constInfo.cmpMaskMode = sparseFlashMLACmpParams.cmpMaskMode;
@@ -372,8 +411,7 @@ __aicore__ inline void SparseFlashMlaCsaKernel<CubeBlockType, VecBlockType>::Ini
372 uint32_t mm1ResultSize = constInfo.s1BaseSize / CV_RATIO * constInfo.s2BaseSize * sizeof(T);411 uint32_t mm1ResultSize = constInfo.s1BaseSize / CV_RATIO * constInfo.s2BaseSize * sizeof(T);
373 uint32_t mm2ResultSize = constInfo.s1BaseSize / CV_RATIO * 512 * sizeof(T);412 uint32_t mm2ResultSize = constInfo.s1BaseSize / CV_RATIO * 512 * sizeof(T);
374 uint32_t mm2LeftSize = constInfo.s1BaseSize * constInfo.s2BaseSize * sizeof(Q_T);413 uint32_t mm2LeftSize = constInfo.s1BaseSize * constInfo.s2BaseSize * sizeof(Q_T);
375- uint32_t mm1RightSize = constInfo.s2BaseSize * 512 * sizeof(Q_T);414+ l1BufferManager.Init(pipe, mm2LeftSize * 2);
376- l1BufferManager.Init(pipe, 524288); // 512 * 1024
377 // 保存p结果的L1内存必须放在第一个L1 policy上，保证和vec申请的地址相同415 // 保存p结果的L1内存必须放在第一个L1 policy上，保证和vec申请的地址相同
378 l1PBuffers.Init(l1BufferManager, mm2LeftSize);416 l1PBuffers.Init(l1BufferManager, mm2LeftSize);
379 l1PBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID);417 l1PBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID);
@@ -381,14 +419,6 @@ __aicore__ inline void SparseFlashMlaCsaKernel<CubeBlockType, VecBlockType>::Ini
381 l1PBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID);419 l1PBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID);
382 crossCoreSyncBufId++;420 crossCoreSyncBufId++;
383 421 
384- l1RightBuffers.Init(l1BufferManager, mm1RightSize);
385- l1RightBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID);
386- crossCoreSyncBufId++;
387- l1RightBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID);
388- crossCoreSyncBufId++;
389- l1RightBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID);
390- crossCoreSyncBufId++;
391- 
392 ubBufferManager.Init(pipe, mm1ResultSize * 2 + mm2ResultSize);422 ubBufferManager.Init(pipe, mm1ResultSize * 2 + mm2ResultSize);
393 bmm2Buffers.Init(ubBufferManager, mm2ResultSize);423 bmm2Buffers.Init(ubBufferManager, mm2ResultSize);
394 bmm2Buffers.Get().SetCrossCoreID(crossCoreSyncBufId, crossCoreSyncBufId);424 bmm2Buffers.Get().SetCrossCoreID(crossCoreSyncBufId, crossCoreSyncBufId);
@@ -438,17 +468,17 @@ __aicore__ inline void SparseFlashMlaCsaKernel<CubeBlockType, VecBlockType>::Ini
438 constexpr uint32_t FD_MAX_SUM_REGION_NUM = 2U;468 constexpr uint32_t FD_MAX_SUM_REGION_NUM = 2U;
439 uint32_t gSize = static_cast<uint32_t>(constInfo.gSize);469 uint32_t gSize = static_cast<uint32_t>(constInfo.gSize);
440 uint32_t combineElemSize = gSize * constInfo.dSize + FD_MAX_SUM_REGION_NUM * gSize *470 uint32_t combineElemSize = gSize * constInfo.dSize + FD_MAX_SUM_REGION_NUM * gSize *
441- static_cast<uint32_t>(AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW);471+ static_cast<uint32_t>(AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW);
442 if constexpr (IS_BATCH_CONSISTENCY) {472 if constexpr (IS_BATCH_CONSISTENCY) {
443 uint32_t intraCoreSlotNum = IS_SPLIT_G ? GetBlockNum() : (GetBlockNum() << 1U);473 uint32_t intraCoreSlotNum = IS_SPLIT_G ? GetBlockNum() : (GetBlockNum() << 1U);
444 uint32_t intraCoreCombineSize = intraCoreSlotNum * combineElemSize * sizeof(float);474 uint32_t intraCoreCombineSize = intraCoreSlotNum * combineElemSize * sizeof(float);
445 uint32_t crossCoreCombineSize = GetBlockNum() * BATCH_CONSISTENCY_MAX_REDUCE_BLOCK_NUM *475 uint32_t crossCoreCombineSize = GetBlockNum() * BATCH_CONSISTENCY_MAX_REDUCE_BLOCK_NUM *
446- combineElemSize * sizeof(float);476+ combineElemSize * sizeof(float);
447 intraCoreCombineBuffer.Init(fdStagingBufferManager, intraCoreCombineSize);477 intraCoreCombineBuffer.Init(fdStagingBufferManager, intraCoreCombineSize);
448 crossCoreCombineBuffer.Init(fdStagingBufferManager, crossCoreCombineSize);478 crossCoreCombineBuffer.Init(fdStagingBufferManager, crossCoreCombineSize);
449 } else {479 } else {
450 uint32_t fdSlotCount = static_cast<uint32_t>(AttentionCommon::FD_MAX_S2_SPLIT_NUM) *480 uint32_t fdSlotCount = static_cast<uint32_t>(AttentionCommon::FD_MAX_S2_SPLIT_NUM) *
451- (IS_SPLIT_G ? (GetBlockNum() >> 1U) : GetBlockNum());481+ (IS_SPLIT_G ? (GetBlockNum() >> 1U) : GetBlockNum());
452 fdStagingBuffer.Init(fdStagingBufferManager, fdSlotCount * combineElemSize * sizeof(float));482 fdStagingBuffer.Init(fdStagingBufferManager, fdSlotCount * combineElemSize * sizeof(float));
453 }483 }
454}484}
@@ -527,11 +557,11 @@ __aicore__ inline void SparseFlashMlaCsaKernel<CubeBlockType, VecBlockType>::Pro
527 maxS2LoopCnt = static_cast<int64_t>(metadataGm.GetValue(GetAttrAbsIndex(aicIdx, FA_S2_MAX_NUM, false)));557 maxS2LoopCnt = static_cast<int64_t>(metadataGm.GetValue(GetAttrAbsIndex(aicIdx, FA_S2_MAX_NUM, false)));
528 }558 }
529 if (hasLoad == 0) {559 if (hasLoad == 0) {
530- if ASCEND_IS_AIV {560+ if ASCEND_IS_AIC {
531 if constexpr (IS_SPLIT_G) {561 if constexpr (IS_SPLIT_G) {
532 for (int64_t loopCnt = 0; loopCnt < maxS2LoopCnt; loopCnt++) {562 for (int64_t loopCnt = 0; loopCnt < maxS2LoopCnt; loopCnt++) {
533- CrossCoreSetFlag<0, PIPE_MTE3>(15);563+ CrossCoreSetFlag<0, PIPE_MTE2>(crossCoreMte2SyncFlagId);
534- CrossCoreWaitFlag<0, PIPE_MTE3>(15);564+ CrossCoreWaitFlag<0, PIPE_MTE2>(crossCoreMte2SyncFlagId);
535 }565 }
536 }566 }
537 }567 }
@@ -556,7 +586,8 @@ __aicore__ inline void SparseFlashMlaCsaKernel<CubeBlockType, VecBlockType>::Pro
556 586 
557 int64_t taskId = 0;587 int64_t taskId = 0;
558 bool notLast = true;588 bool notLast = true;
559- RunInfo runInfo[3];589+ bool isFirstLoop = true;
590+ RunInfo runInfo[4];
560 RunParamStr runParam;591 RunParamStr runParam;
561 runParam.firstFdDataWorkspaceIdx = firstFdDataWorkspaceIdx;592 runParam.firstFdDataWorkspaceIdx = firstFdDataWorkspaceIdx;
562 int64_t multiCoreInnerIdx = 1;593 int64_t multiCoreInnerIdx = 1;
@@ -573,22 +604,28 @@ __aicore__ inline void SparseFlashMlaCsaKernel<CubeBlockType, VecBlockType>::Pro
573 604 
574 int64_t gS1LoopEnd = lastBN ? (runParam.gs1LoopEndIdx + PRELOAD_NUM) : runParam.gs1LoopEndIdx;605 int64_t gS1LoopEnd = lastBN ? (runParam.gs1LoopEndIdx + PRELOAD_NUM) : runParam.gs1LoopEndIdx;
575 for (int64_t gS1Index = runParam.gs1LoopStartIdx; gS1Index < gS1LoopEnd; gS1Index++) {606 for (int64_t gS1Index = runParam.gs1LoopStartIdx; gS1Index < gS1LoopEnd; gS1Index++) {
607+ bool notLastThreeLoop = true;
576 bool notLastTwoLoop = true;608 bool notLastTwoLoop = true;
577 if (lastBN) {609 if (lastBN) {
578 int32_t extraGS1 = gS1Index - runParam.gs1LoopEndIdx;610 int32_t extraGS1 = gS1Index - runParam.gs1LoopEndIdx;
579 switch (extraGS1) {611 switch (extraGS1) {
580 case 0:612 case 0:
581- notLastTwoLoop = false;613+ notLastThreeLoop = false;
582 break;614 break;
583 case 1:615 case 1:
616+ notLastTwoLoop = false;
617+ notLastThreeLoop = false;
618+ break;
619+ case 2:
584 notLast = false;620 notLast = false;
585 notLastTwoLoop = false;621 notLastTwoLoop = false;
622+ notLastThreeLoop = false;
586 break;623 break;
587 default:624 default:
588 break;625 break;
589 }626 }
590 }627 }
591- if (notLastTwoLoop) {628+ if (notLastThreeLoop) {
592 this->ComputeAxisIdxByBnAndGs1(bnIdx, gS1Index, runParam);629 this->ComputeAxisIdxByBnAndGs1(bnIdx, gS1Index, runParam);
593 bool s1NoNeedCalc =630 bool s1NoNeedCalc =
594 ComputeParamS1<TEMPLATE_INTF_ARGS>(runParam, this->constInfo, gS1Index, this->cuSeqlensQGm);631 ComputeParamS1<TEMPLATE_INTF_ARGS>(runParam, this->constInfo, gS1Index, this->cuSeqlensQGm);
@@ -637,7 +674,8 @@ __aicore__ inline void SparseFlashMlaCsaKernel<CubeBlockType, VecBlockType>::Pro
637 for (int64_t s2LoopCount = 0; s2LoopCount <= s2LoopLimit; ++s2LoopCount) {674 for (int64_t s2LoopCount = 0; s2LoopCount <= s2LoopLimit; ++s2LoopCount) {
638 if constexpr (IS_BATCH_CONSISTENCY) {675 if constexpr (IS_BATCH_CONSISTENCY) {
639 int64_t safeBaseBlockNum = runParam.baseBlockNumPerReductionBlock > 0 ?676 int64_t safeBaseBlockNum = runParam.baseBlockNumPerReductionBlock > 0 ?
640- runParam.baseBlockNumPerReductionBlock : 1LL;677+ runParam.baseBlockNumPerReductionBlock :
678+ 1LL;
641 if (runParam.isCrossCoreSplit && s2LoopCount % safeBaseBlockNum == 0) {679 if (runParam.isCrossCoreSplit && s2LoopCount % safeBaseBlockNum == 0) {
642 runParam.s2SplitIdx = s2SplitIdxCounter++;680 runParam.s2SplitIdx = s2SplitIdxCounter++;
643 }681 }
@@ -645,42 +683,50 @@ __aicore__ inline void SparseFlashMlaCsaKernel<CubeBlockType, VecBlockType>::Pro
645 if constexpr (TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE ||683 if constexpr (TEMPLATE_MODE == SMLATemplateMode::CSA_TEMPLATE_MODE ||
646 TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||684 TEMPLATE_MODE == SMLATemplateMode::ORI_SPARSE_TEMPLATE_MODE ||
647 TEMPLATE_MODE == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) {685 TEMPLATE_MODE == SMLATemplateMode::ORI_CMP_SPARSE_TEMPLATE_MODE) {
648- if (notLastTwoLoop) {686+ if (notLastThreeLoop) {
649- RunInfo &runInfo1 = runInfo[taskId % 3];687+ RunInfo &runInfo1 = runInfo[taskId % 4];
650 this->SetRunInfo(runInfo1, runParam, taskId, s2LoopCount, s2LoopLimit, multiCoreInnerIdx);688 this->SetRunInfo(runInfo1, runParam, taskId, s2LoopCount, s2LoopLimit, multiCoreInnerIdx);
651- if ASCEND_IS_AIC {689+ }
652- this->cubeBlock.IterateBmm1(this->bmm1Buffers.Get(), this->l1RightBuffers.Get(),690+ if ASCEND_IS_AIV {
653- v0ResGmBuffers.Get(), runInfo1, this->constInfo);691+ if (notLastThreeLoop) {
654- } else {692+ RunInfo &runInfo1 = runInfo[taskId % 4];
655- this->vecBlock.ProcessVec0(this->l1RightBuffers.Get(), v0ResGmBuffers.Get(), runInfo1,693+ this->vecBlock.ProcessVec0(this->v0ResGmBuffers.Get(runInfo1.taskIdMod3), runInfo1,
656- this->constInfo, 0);694+ this->constInfo);
695+ }
696+ if (taskId > 1 && notLast) {
697+ auto &runInfo2 = runInfo[(taskId + 2) % 4];
698+ this->vecBlock.ProcessVec1(this->l1PBuffers.Get(), this->bmm1Buffers.Get(), runInfo2,
699+ this->constInfo);
700+ }
701+ if (taskId > 2) {
702+ RunInfo &runInfo3 = runInfo[(taskId + 1) % 4];
703+ this->vecBlock.ProcessVec2(this->bmm2Buffers.Get(), runInfo3, this->constInfo);
657 }704 }
658 } else {705 } else {
659- if ASCEND_IS_AIV {706+ if (taskId > 0 && notLastTwoLoop) {
707+ RunInfo &runInfo1 = runInfo[(taskId + 3) % 4];
708+ this->cubeBlock.IterateLoadQK(this->v0ResGmBuffers.Get(runInfo1.taskIdMod3),
709+ runInfo1, this->constInfo, isFirstLoop);
710+ isFirstLoop = false;
711+ } else {
660 if constexpr (IS_SPLIT_G) {712 if constexpr (IS_SPLIT_G) {
661- if (maxS2LoopCnt > 0) {713+ if (taskId > 0 && maxS2LoopCnt > 0) {
662 maxS2LoopCnt--;714 maxS2LoopCnt--;
663- CrossCoreSetFlag<0, PIPE_MTE3>(15);715+ CrossCoreSetFlag<0, PIPE_MTE2>(crossCoreMte2SyncFlagId);
664- CrossCoreWaitFlag<0, PIPE_MTE3>(15);716+ CrossCoreWaitFlag<0, PIPE_MTE2>(crossCoreMte2SyncFlagId);
665 }717 }
666 }718 }
667 }719 }
668- }720+ if (taskId > 1 && notLast) {
669- if (taskId > 0 && notLast) {721+ auto &runInfo2 = runInfo[(taskId + 2) % 4];
670- auto &runInfo2 = runInfo[(taskId + 2) % 3];722+ RunInfo &runInfoNext = runInfo[(taskId + 3) % 4];
671- if ASCEND_IS_AIV {723+ this->cubeBlock.IterateBmm1(this->bmm1Buffers.Get(),
672- this->vecBlock.ProcessVec1(this->l1PBuffers.Get(), this->bmm1Buffers.Get(), runInfo2,724+ notLastTwoLoop, runInfoNext, runInfo2, this->constInfo);
673- this->constInfo);
674- } else {
675- RunInfo &runInfo2 = runInfo[(taskId + 2) % 3];
676- this->cubeBlock.IterateBmm2(this->bmm2Buffers.Get(), this->l1PBuffers,
677- this->l1RightBuffers.GetReused(), runInfo2, this->constInfo);
678 }725 }
679- }726+ if (taskId > 2) {
680- if (taskId > 1) {727+ RunInfo &runInfo3 = runInfo[(taskId + 1) % 4];
681- if ASCEND_IS_AIV {728+ this->cubeBlock.IterateBmm2(this->bmm2Buffers.Get(), this->l1PBuffers,
682- RunInfo &runInfo3 = runInfo[(taskId + 1) % 3];729+ runInfo3, this->constInfo);
683- this->vecBlock.ProcessVec2(this->bmm2Buffers.Get(), runInfo3, this->constInfo);
684 }730 }
685 }731 }
686 }732 }
@@ -690,11 +736,11 @@ __aicore__ inline void SparseFlashMlaCsaKernel<CubeBlockType, VecBlockType>::Pro
690 }736 }
691 gS1StartIdx = 0;737 gS1StartIdx = 0;
692 }738 }
693- if ASCEND_IS_AIV {739+ if ASCEND_IS_AIC {
694 if constexpr (IS_SPLIT_G) {740 if constexpr (IS_SPLIT_G) {
695 for (int64_t loopCnt = 0; loopCnt < maxS2LoopCnt; loopCnt++) {741 for (int64_t loopCnt = 0; loopCnt < maxS2LoopCnt; loopCnt++) {
696- CrossCoreSetFlag<0, PIPE_MTE3>(15);742+ CrossCoreSetFlag<0, PIPE_MTE2>(crossCoreMte2SyncFlagId);
697- CrossCoreWaitFlag<0, PIPE_MTE3>(15);743+ CrossCoreWaitFlag<0, PIPE_MTE2>(crossCoreMte2SyncFlagId);
698 }744 }
699 }745 }
700 }746 }
@@ -752,12 +798,13 @@ __aicore__ inline void SparseFlashMlaCsaKernel<CubeBlockType, VecBlockType>::Set
752 runInfo.s2SplitIdx = runParam.s2SplitIdx;798 runInfo.s2SplitIdx = runParam.s2SplitIdx;
753 runInfo.isFirstS2SplitCore = runParam.isFirstS2SplitCore;799 runInfo.isFirstS2SplitCore = runParam.isFirstS2SplitCore;
754 int64_t safeBaseBlockNum = runParam.baseBlockNumPerReductionBlock > 0 ?800 int64_t safeBaseBlockNum = runParam.baseBlockNumPerReductionBlock > 0 ?
755- runParam.baseBlockNumPerReductionBlock : 1LL;801+ runParam.baseBlockNumPerReductionBlock :
802+ 1LL;
756 int64_t baseBlockIdInReduceBlock = s2LoopCount % safeBaseBlockNum;803 int64_t baseBlockIdInReduceBlock = s2LoopCount % safeBaseBlockNum;
757 runInfo.reduceBlockId = s2LoopCount / safeBaseBlockNum;804 runInfo.reduceBlockId = s2LoopCount / safeBaseBlockNum;
758 runInfo.isFirstBase = baseBlockIdInReduceBlock == 0;805 runInfo.isFirstBase = baseBlockIdInReduceBlock == 0;
759 runInfo.isLastBase = baseBlockIdInReduceBlock == safeBaseBlockNum - 1LL ||806 runInfo.isLastBase = baseBlockIdInReduceBlock == safeBaseBlockNum - 1LL ||
760- s2LoopCount == s2LoopLimit;807+ s2LoopCount == s2LoopLimit;
761 runInfo.needReduce = runInfo.reduceBlockId > 0;808 runInfo.needReduce = runInfo.reduceBlockId > 0;
762 this->ComputeBmm1Tail(runInfo, runParam);809 this->ComputeBmm1Tail(runInfo, runParam);
763 InitUniqueRunInfo(runParam, runInfo);810 InitUniqueRunInfo(runParam, runInfo);
Mattention/sparse_flash_mla/op_kernel/arch35/sparse_flash_mla_swa_kernel.h+23-17
@@ -27,6 +27,9 @@
27#else27#else
28#include "../common/matmul.h"28#include "../common/matmul.h"
29#endif29#endif
30+#if __has_include("common/buffers_policy_3buff_sfa.h")
31+#include "common/buffers_policy_3buff_sfa.h"
32+#endif
30#if __has_include("../../common/op_kernel/FixpipeOut.h")33#if __has_include("../../common/op_kernel/FixpipeOut.h")
31#include "../../common/op_kernel/FixpipeOut.h"34#include "../../common/op_kernel/FixpipeOut.h"
32#else35#else
@@ -116,7 +119,6 @@ private:
116 // mm2左矩阵P119 // mm2左矩阵P
117 BufferManager<BufferType::L1> l1BufferManager;120 BufferManager<BufferType::L1> l1BufferManager;
118 BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> l1PBuffers;121 BuffersPolicyDB<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> l1PBuffers;
119- BuffersPolicy3buff<BufferType::L1, SyncType::INNER_CORE_SYNC> l1RightBuffers;
120 /* GM信息 */122 /* GM信息 */
121 GlobalTensor<uint32_t> metadataGm;123 GlobalTensor<uint32_t> metadataGm;
122 GlobalTensor<int32_t> cuSeqlensQGm;124 GlobalTensor<int32_t> cuSeqlensQGm;
@@ -191,7 +193,7 @@ __aicore__ inline void SparseFlashMlaSwaKernel<CubeBlockType, VecBlockType>::Ini
191 } else {193 } else {
192 vecBlock.InitS2SplitStaging(fdStagingBuffer.Get());194 vecBlock.InitS2SplitStaging(fdStagingBuffer.Get());
193 }195 }
194- cubeBlock.InitCubeBlock(pipe, l1BufferManager, query);196+ cubeBlock.InitCubeBlock(pipe, query);
195 this->ComputeConstexpr();197 this->ComputeConstexpr();
196 this->InitGlobalBuffer(query, oriKV, cmpKV, oriSparseIndices, cmpSparseIndices, oriBlockTable, cmpBlockTable,198 this->InitGlobalBuffer(query, oriKV, cmpKV, oriSparseIndices, cmpSparseIndices, oriBlockTable, cmpBlockTable,
197 cuSeqlensQ, cuSeqlensOriKv, cuSeqlensCmpKv, sequsedQ, seqUsedOriKv, seqUsedCmpKv,199 cuSeqlensQ, cuSeqlensOriKv, cuSeqlensCmpKv, sequsedQ, seqUsedOriKv, seqUsedCmpKv,
@@ -344,8 +346,7 @@ __aicore__ inline void SparseFlashMlaSwaKernel<CubeBlockType, VecBlockType>::Ini
344 uint32_t mm1ResultSize = constInfo.s1BaseSize / CV_RATIO * constInfo.s2BaseSize * sizeof(T);346 uint32_t mm1ResultSize = constInfo.s1BaseSize / CV_RATIO * constInfo.s2BaseSize * sizeof(T);
345 uint32_t mm2ResultSize = constInfo.s1BaseSize / CV_RATIO * 512 * sizeof(T);347 uint32_t mm2ResultSize = constInfo.s1BaseSize / CV_RATIO * 512 * sizeof(T);
346 uint32_t mm2LeftSize = constInfo.s1BaseSize * constInfo.s2BaseSize * sizeof(Q_T);348 uint32_t mm2LeftSize = constInfo.s1BaseSize * constInfo.s2BaseSize * sizeof(Q_T);
347- uint32_t mm1RightSize = constInfo.s2BaseSize * 512 * sizeof(Q_T);349+ l1BufferManager.Init(pipe, mm2LeftSize * 2);
348- l1BufferManager.Init(pipe, 524288); // 512 * 1024
349 // 保存p结果的L1内存必须放在第一个L1 policy上，保证和vec申请的地址相同350 // 保存p结果的L1内存必须放在第一个L1 policy上，保证和vec申请的地址相同
350 l1PBuffers.Init(l1BufferManager, mm2LeftSize);351 l1PBuffers.Init(l1BufferManager, mm2LeftSize);
351 l1PBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID);352 l1PBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID);
@@ -353,7 +354,6 @@ __aicore__ inline void SparseFlashMlaSwaKernel<CubeBlockType, VecBlockType>::Ini
353 l1PBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID);354 l1PBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID);
354 crossCoreSyncBufId++;355 crossCoreSyncBufId++;
355 if ASCEND_IS_AIC {356 if ASCEND_IS_AIC {
356- l1RightBuffers.Init(l1BufferManager, mm1RightSize);
357 l1PBuffers.Get().SetCrossCore();357 l1PBuffers.Get().SetCrossCore();
358 l1PBuffers.Get().SetCrossCore();358 l1PBuffers.Get().SetCrossCore();
359 }359 }
@@ -385,17 +385,17 @@ __aicore__ inline void SparseFlashMlaSwaKernel<CubeBlockType, VecBlockType>::Ini
385 constexpr uint32_t FD_MAX_SUM_REGION_NUM = 2U;385 constexpr uint32_t FD_MAX_SUM_REGION_NUM = 2U;
386 uint32_t gSize = static_cast<uint32_t>(constInfo.gSize);386 uint32_t gSize = static_cast<uint32_t>(constInfo.gSize);
387 uint32_t combineElemSize = gSize * constInfo.dSize + FD_MAX_SUM_REGION_NUM * gSize *387 uint32_t combineElemSize = gSize * constInfo.dSize + FD_MAX_SUM_REGION_NUM * gSize *
388- static_cast<uint32_t>(AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW);388+ static_cast<uint32_t>(AttentionCommon::FD_BROADCAST_ELEMS_PER_ROW);
389 if constexpr (IS_BATCH_CONSISTENCY) {389 if constexpr (IS_BATCH_CONSISTENCY) {
390 uint32_t intraCoreSlotNum = IS_SPLIT_G ? GetBlockNum() : (GetBlockNum() << 1U);390 uint32_t intraCoreSlotNum = IS_SPLIT_G ? GetBlockNum() : (GetBlockNum() << 1U);
391 uint32_t intraCoreCombineSize = intraCoreSlotNum * combineElemSize * sizeof(float);391 uint32_t intraCoreCombineSize = intraCoreSlotNum * combineElemSize * sizeof(float);
392 uint32_t crossCoreCombineSize = GetBlockNum() * BATCH_CONSISTENCY_MAX_REDUCE_BLOCK_NUM *392 uint32_t crossCoreCombineSize = GetBlockNum() * BATCH_CONSISTENCY_MAX_REDUCE_BLOCK_NUM *
393- combineElemSize * sizeof(float);393+ combineElemSize * sizeof(float);
394 intraCoreCombineBuffer.Init(fdStagingBufferManager, intraCoreCombineSize);394 intraCoreCombineBuffer.Init(fdStagingBufferManager, intraCoreCombineSize);
395 crossCoreCombineBuffer.Init(fdStagingBufferManager, crossCoreCombineSize);395 crossCoreCombineBuffer.Init(fdStagingBufferManager, crossCoreCombineSize);
396 } else {396 } else {
397 uint32_t fdSlotCount = static_cast<uint32_t>(AttentionCommon::FD_MAX_S2_SPLIT_NUM) *397 uint32_t fdSlotCount = static_cast<uint32_t>(AttentionCommon::FD_MAX_S2_SPLIT_NUM) *
398- (IS_SPLIT_G ? (GetBlockNum() >> 1U) : GetBlockNum());398+ (IS_SPLIT_G ? (GetBlockNum() >> 1U) : GetBlockNum());
399 fdStagingBuffer.Init(fdStagingBufferManager, fdSlotCount * combineElemSize * sizeof(float));399 fdStagingBuffer.Init(fdStagingBufferManager, fdSlotCount * combineElemSize * sizeof(float));
400 }400 }
401}401}
@@ -487,6 +487,7 @@ __aicore__ inline void SparseFlashMlaSwaKernel<CubeBlockType, VecBlockType>::Pro
487 487 
488 int64_t taskId = 0;488 int64_t taskId = 0;
489 bool notLast = true;489 bool notLast = true;
490+ bool isFirstLoop = true;
490 RunInfo runInfo[3];491 RunInfo runInfo[3];
491 RunParamStr runParam;492 RunParamStr runParam;
492 runParam.firstFdDataWorkspaceIdx = firstFdDataWorkspaceIdx;493 runParam.firstFdDataWorkspaceIdx = firstFdDataWorkspaceIdx;
@@ -566,7 +567,8 @@ __aicore__ inline void SparseFlashMlaSwaKernel<CubeBlockType, VecBlockType>::Pro
566 for (int64_t s2LoopCount = 0; s2LoopCount <= s2LoopLimit; ++s2LoopCount) {567 for (int64_t s2LoopCount = 0; s2LoopCount <= s2LoopLimit; ++s2LoopCount) {
567 if constexpr (IS_BATCH_CONSISTENCY) {568 if constexpr (IS_BATCH_CONSISTENCY) {
568 int64_t safeBaseBlockNum = runParam.baseBlockNumPerReductionBlock > 0 ?569 int64_t safeBaseBlockNum = runParam.baseBlockNumPerReductionBlock > 0 ?
569- runParam.baseBlockNumPerReductionBlock : 1LL;570+ runParam.baseBlockNumPerReductionBlock :
571+ 1LL;
570 if (runParam.isCrossCoreSplit && s2LoopCount % safeBaseBlockNum == 0) {572 if (runParam.isCrossCoreSplit && s2LoopCount % safeBaseBlockNum == 0) {
571 runParam.s2SplitIdx = s2SplitIdxCounter++;573 runParam.s2SplitIdx = s2SplitIdxCounter++;
572 }574 }
@@ -575,8 +577,8 @@ __aicore__ inline void SparseFlashMlaSwaKernel<CubeBlockType, VecBlockType>::Pro
575 RunInfo &runInfo1 = runInfo[taskId % 3];577 RunInfo &runInfo1 = runInfo[taskId % 3];
576 this->SetRunInfo(runInfo1, runParam, taskId, s2LoopCount, s2LoopLimit, multiCoreInnerIdx);578 this->SetRunInfo(runInfo1, runParam, taskId, s2LoopCount, s2LoopLimit, multiCoreInnerIdx);
577 if ASCEND_IS_AIC {579 if ASCEND_IS_AIC {
578- this->cubeBlock.IterateBmm1(this->bmm1Buffers.Get(), this->l1RightBuffers.Get(), runInfo1,580+ this->cubeBlock.IterateLoadQK(runInfo1, this->constInfo, isFirstLoop);
579- this->constInfo);581+ isFirstLoop = false;
580 }582 }
581 }583 }
582 if (taskId > 0 && notLast) {584 if (taskId > 0 && notLast) {
@@ -585,15 +587,18 @@ __aicore__ inline void SparseFlashMlaSwaKernel<CubeBlockType, VecBlockType>::Pro
585 this->vecBlock.ProcessVec1(this->l1PBuffers.Get(), this->bmm1Buffers.Get(), runInfo2,587 this->vecBlock.ProcessVec1(this->l1PBuffers.Get(), this->bmm1Buffers.Get(), runInfo2,
586 this->constInfo);588 this->constInfo);
587 } else {589 } else {
588- RunInfo &runInfo2 = runInfo[(taskId + 2) % 3];590+ RunInfo &runInfoNext = runInfo[taskId % 3];
589- this->cubeBlock.IterateBmm2(this->bmm2Buffers.Get(), this->l1PBuffers,591+ this->cubeBlock.IterateBmm1(this->bmm1Buffers.Get(),
590- this->l1RightBuffers.GetReused(), runInfo2, this->constInfo);592+ notLastTwoLoop, runInfoNext, runInfo2, this->constInfo);
591 }593 }
592 }594 }
593 if (taskId > 1) {595 if (taskId > 1) {
596+ RunInfo &runInfo3 = runInfo[(taskId + 1) % 3];
594 if ASCEND_IS_AIV {597 if ASCEND_IS_AIV {
595- RunInfo &runInfo3 = runInfo[(taskId + 1) % 3];
596 this->vecBlock.ProcessVec2(this->bmm2Buffers.Get(), runInfo3, this->constInfo);598 this->vecBlock.ProcessVec2(this->bmm2Buffers.Get(), runInfo3, this->constInfo);
599+ } else {
600+ this->cubeBlock.IterateBmm2(this->bmm2Buffers.Get(), this->l1PBuffers,
601+ runInfo3, this->constInfo);
597 }602 }
598 }603 }
599 ++taskId;604 ++taskId;
@@ -656,12 +661,13 @@ __aicore__ inline void SparseFlashMlaSwaKernel<CubeBlockType, VecBlockType>::Set
656 runInfo.s2SplitIdx = runParam.s2SplitIdx;661 runInfo.s2SplitIdx = runParam.s2SplitIdx;
657 runInfo.isFirstS2SplitCore = runParam.isFirstS2SplitCore;662 runInfo.isFirstS2SplitCore = runParam.isFirstS2SplitCore;
658 int64_t safeBaseBlockNum = runParam.baseBlockNumPerReductionBlock > 0 ?663 int64_t safeBaseBlockNum = runParam.baseBlockNumPerReductionBlock > 0 ?
659- runParam.baseBlockNumPerReductionBlock : 1LL;664+ runParam.baseBlockNumPerReductionBlock :
665+ 1LL;
660 int64_t baseBlockIdInReduceBlock = s2LoopCount % safeBaseBlockNum;666 int64_t baseBlockIdInReduceBlock = s2LoopCount % safeBaseBlockNum;
661 runInfo.reduceBlockId = s2LoopCount / safeBaseBlockNum;667 runInfo.reduceBlockId = s2LoopCount / safeBaseBlockNum;
662 runInfo.isFirstBase = baseBlockIdInReduceBlock == 0;668 runInfo.isFirstBase = baseBlockIdInReduceBlock == 0;
663 runInfo.isLastBase = baseBlockIdInReduceBlock == safeBaseBlockNum - 1LL ||669 runInfo.isLastBase = baseBlockIdInReduceBlock == safeBaseBlockNum - 1LL ||
664- s2LoopCount == s2LoopLimit;670+ s2LoopCount == s2LoopLimit;
665 runInfo.needReduce = runInfo.reduceBlockId > 0;671 runInfo.needReduce = runInfo.reduceBlockId > 0;
666 this->ComputeBmm1Tail(runInfo, runParam);672 this->ComputeBmm1Tail(runInfo, runParam);
667 InitUniqueRunInfo(runParam, runInfo);673 InitUniqueRunInfo(runParam, runInfo);
Mattention/sparse_flash_mla/op_kernel/arch35/util_regbase.h+174-172文件内容审核中，请稍后刷新重试
Mattention/sparse_flash_mla/op_kernel/sparse_flash_mla.cpp+29-29
@@ -34,37 +34,37 @@ using namespace optiling::detail;
34using namespace SMLAKernel;34using namespace SMLAKernel;
35 35 
36#if (__CCE_AICORE__ == 310)36#if (__CCE_AICORE__ == 310)
37-#define SMLA_OP_IMPL(templateClass, tilingdataClass, ...) \37+#define SMLA_OP_IMPL(templateClass, tilingdataClass, ...) \
38- do { \38+ do { \
39- using CubeBlockType = \39+ using CubeBlockType = \
40- typename std::conditional<g_coreType == AscendC::AIC, SMLAKernel::CSABlockCube<__VA_ARGS__>, \40+ typename std::conditional<g_coreType == AscendC::AIC, SMLAKernel::CSABlockCube<__VA_ARGS__>, \
41- SMLAKernel::CSABlockCubeDummy<__VA_ARGS__>>::type; \41+ SMLAKernel::CSABlockCubeDummy<__VA_ARGS__>>::type; \
42- using VecBlockType = \42+ using VecBlockType = \
43- typename std::conditional<g_coreType == AscendC::AIC, SMLAKernel::CSABlockVecDummy<__VA_ARGS__>, \43+ typename std::conditional<g_coreType == AscendC::AIC, SMLAKernel::CSABlockVecDummy<__VA_ARGS__>, \
44- SMLAKernel::CSABlockVec<__VA_ARGS__>>::type; \44+ SMLAKernel::CSABlockVec<__VA_ARGS__>>::type; \
45- templateClass<CubeBlockType, VecBlockType> op; \45+ templateClass<CubeBlockType, VecBlockType> op; \
46- GET_TILING_DATA_WITH_STRUCT(tilingdataClass, tilingDataIn, tiling); \46+ GET_TILING_DATA_WITH_STRUCT(tilingdataClass, tilingDataIn, tiling); \
47- const tilingdataClass *__restrict tilingData = &tilingDataIn; \47+ const tilingdataClass *__restrict tilingData = &tilingDataIn; \
48- op.Init(query, oriKV, cmpKV, oriSparseIndices, cmpSparseIndices, oriBlockTable, cmpBlockTable, cuSeqlensQ, \48+ op.Init(query, oriKV, cmpKV, oriSparseIndices, cmpSparseIndices, oriBlockTable, cmpBlockTable, cuSeqlensQ, \
49- cuSeqlensOriKv, cuSeqlensCmpKv, seqUsedQ, seqUsedOriKV, seqUsedCmpKV, cmpResidualKV, oriTopkLength, \49+ cuSeqlensOriKv, cuSeqlensCmpKv, seqUsedQ, seqUsedOriKV, seqUsedCmpKV, cmpResidualKV, oriTopkLength, \
50- cmpTopkLength, sinks, metadata, attentionOut, softmaxLse, user, tilingData, &tPipe); \50+ cmpTopkLength, sinks, metadata, attentionOut, softmaxLse, user, tilingData, &tPipe); \
51- op.Process(); \51+ op.Process(); \
52 } while (0)52 } while (0)
53#else53#else
54-#define SMLA_OP_IMPL(templateClass, tilingdataClass, ...) \54+#define SMLA_OP_IMPL(templateClass, tilingdataClass, ...) \
55- do { \55+ do { \
56- templateClass<SMLAType<__VA_ARGS__>> op; \56+ templateClass<SMLAType<__VA_ARGS__>> op; \
57- GET_TILING_DATA_WITH_STRUCT(tilingdataClass, tiling_data_in, tiling); \57+ GET_TILING_DATA_WITH_STRUCT(tilingdataClass, tiling_data_in, tiling); \
58- const tilingdataClass *__restrict tiling_data = &tiling_data_in; \58+ const tilingdataClass *__restrict tiling_data = &tiling_data_in; \
59- op.Init(query, oriKV, cmpKV, cmpSparseIndices, oriBlockTable, cmpBlockTable, cuSeqlensQ, cuSeqlensOriKv, \59+ op.Init(query, oriKV, cmpKV, cmpSparseIndices, oriBlockTable, cmpBlockTable, cuSeqlensQ, cuSeqlensOriKv, \
60- cuSeqlensCmpKv, seqUsedQ, seqUsedOriKV, seqUsedCmpKV, cmpResidualKV, sinks, metadata, attentionOut, \60+ cuSeqlensCmpKv, seqUsedQ, seqUsedOriKV, seqUsedCmpKV, cmpResidualKV, sinks, metadata, attentionOut, \
61- softmaxLse, user, tiling_data, tiling, &tPipe); \61+ softmaxLse, user, tiling_data, tiling, &tPipe); \
62- op.Process(); \62+ op.Process(); \
63 } while (0)63 } while (0)
64#endif64#endif
65 65 
66template <int FLASH_DECODE, int LAYOUT_T, int KV_LAYOUT_T, int TEMPLATE_MODE, int SPLIT_G, int HEAD_RATIO_ONE,66template <int FLASH_DECODE, int LAYOUT_T, int KV_LAYOUT_T, int TEMPLATE_MODE, int SPLIT_G, int HEAD_RATIO_ONE,
67- int BATCH_CONSISTENCY>67+ int BATCH_CONSISTENCY, int IS_VEC_S2PHYADDR>
68__global__ __aicore__ void68__global__ __aicore__ void
69sparse_flash_mla(__gm__ uint8_t *query, __gm__ uint8_t *oriKV, __gm__ uint8_t *cmpKV, __gm__ uint8_t *oriSparseIndices,69sparse_flash_mla(__gm__ uint8_t *query, __gm__ uint8_t *oriKV, __gm__ uint8_t *cmpKV, __gm__ uint8_t *oriSparseIndices,
70 __gm__ uint8_t *cmpSparseIndices, __gm__ uint8_t *oriBlockTable, __gm__ uint8_t *cmpBlockTable,70 __gm__ uint8_t *cmpSparseIndices, __gm__ uint8_t *oriBlockTable, __gm__ uint8_t *cmpBlockTable,
@@ -85,11 +85,11 @@ sparse_flash_mla(__gm__ uint8_t *query, __gm__ uint8_t *oriKV, __gm__ uint8_t *c
85 TEMPLATE_MODE == ORI_CMP_SPARSE_TEMPLATE) {85 TEMPLATE_MODE == ORI_CMP_SPARSE_TEMPLATE) {
86 SMLA_OP_IMPL(SMLAKernel::SparseFlashMlaCsaKernel, SparseFlashMlaTilingData, half, half, float, half,86 SMLA_OP_IMPL(SMLAKernel::SparseFlashMlaCsaKernel, SparseFlashMlaTilingData, half, half, float, half,
87 FLASH_DECODE, static_cast<SMLA_LAYOUT>(LAYOUT_T), static_cast<SMLA_LAYOUT>(KV_LAYOUT_T),87 FLASH_DECODE, static_cast<SMLA_LAYOUT>(LAYOUT_T), static_cast<SMLA_LAYOUT>(KV_LAYOUT_T),
88- static_cast<SMLATemplateMode>(TEMPLATE_MODE), SPLIT_G, BATCH_CONSISTENCY);88+ static_cast<SMLATemplateMode>(TEMPLATE_MODE), SPLIT_G, BATCH_CONSISTENCY, IS_VEC_S2PHYADDR);
89 } else {89 } else {
90 SMLA_OP_IMPL(SMLAKernel::SparseFlashMlaSwaKernel, SparseFlashMlaTilingData, half, half, float, half,90 SMLA_OP_IMPL(SMLAKernel::SparseFlashMlaSwaKernel, SparseFlashMlaTilingData, half, half, float, half,
91 FLASH_DECODE, static_cast<SMLA_LAYOUT>(LAYOUT_T), static_cast<SMLA_LAYOUT>(KV_LAYOUT_T),91 FLASH_DECODE, static_cast<SMLA_LAYOUT>(LAYOUT_T), static_cast<SMLA_LAYOUT>(KV_LAYOUT_T),
92- static_cast<SMLATemplateMode>(TEMPLATE_MODE), SPLIT_G, BATCH_CONSISTENCY);92+ static_cast<SMLATemplateMode>(TEMPLATE_MODE), SPLIT_G, BATCH_CONSISTENCY, IS_VEC_S2PHYADDR);
93 }93 }
94 }94 }
95 if constexpr (ORIG_DTYPE_Q == DT_BF16 && ORIG_DTYPE_ORI_KV == DT_BF16 && ORIG_DTYPE_ATTN_OUT == DT_BF16) {95 if constexpr (ORIG_DTYPE_Q == DT_BF16 && ORIG_DTYPE_ORI_KV == DT_BF16 && ORIG_DTYPE_ATTN_OUT == DT_BF16) {
@@ -98,12 +98,12 @@ sparse_flash_mla(__gm__ uint8_t *query, __gm__ uint8_t *oriKV, __gm__ uint8_t *c
98 SMLA_OP_IMPL(SMLAKernel::SparseFlashMlaCsaKernel, SparseFlashMlaTilingData, bfloat16_t, bfloat16_t, float,98 SMLA_OP_IMPL(SMLAKernel::SparseFlashMlaCsaKernel, SparseFlashMlaTilingData, bfloat16_t, bfloat16_t, float,
99 bfloat16_t, FLASH_DECODE, static_cast<SMLA_LAYOUT>(LAYOUT_T),99 bfloat16_t, FLASH_DECODE, static_cast<SMLA_LAYOUT>(LAYOUT_T),
100 static_cast<SMLA_LAYOUT>(KV_LAYOUT_T), static_cast<SMLATemplateMode>(TEMPLATE_MODE), SPLIT_G,100 static_cast<SMLA_LAYOUT>(KV_LAYOUT_T), static_cast<SMLATemplateMode>(TEMPLATE_MODE), SPLIT_G,
101- BATCH_CONSISTENCY);101+ BATCH_CONSISTENCY, IS_VEC_S2PHYADDR);
102 } else {102 } else {
103 SMLA_OP_IMPL(SMLAKernel::SparseFlashMlaSwaKernel, SparseFlashMlaTilingData, bfloat16_t, bfloat16_t, float,103 SMLA_OP_IMPL(SMLAKernel::SparseFlashMlaSwaKernel, SparseFlashMlaTilingData, bfloat16_t, bfloat16_t, float,
104 bfloat16_t, FLASH_DECODE, static_cast<SMLA_LAYOUT>(LAYOUT_T),104 bfloat16_t, FLASH_DECODE, static_cast<SMLA_LAYOUT>(LAYOUT_T),
105 static_cast<SMLA_LAYOUT>(KV_LAYOUT_T), static_cast<SMLATemplateMode>(TEMPLATE_MODE), SPLIT_G,105 static_cast<SMLA_LAYOUT>(KV_LAYOUT_T), static_cast<SMLATemplateMode>(TEMPLATE_MODE), SPLIT_G,
106- BATCH_CONSISTENCY);106+ BATCH_CONSISTENCY, IS_VEC_S2PHYADDR);
107 }107 }
108 }108 }
109#else109#else
Mattention/sparse_flash_mla/op_kernel/sparse_flash_mla_template_tiling_key.h+6-3
@@ -40,7 +40,8 @@ ASCENDC_TPL_ARGS_DECL(SparseFlashMla, // 算子OpType
40 HCA_TEMPLATE, CSA_TEMPLATE, ORI_SPARSE_TEMPLATE, ORI_CMP_SPARSE_TEMPLATE),40 HCA_TEMPLATE, CSA_TEMPLATE, ORI_SPARSE_TEMPLATE, ORI_CMP_SPARSE_TEMPLATE),
41 ASCENDC_TPL_BOOL_DECL(SPLIT_G, 0, 1),41 ASCENDC_TPL_BOOL_DECL(SPLIT_G, 0, 1),
42 ASCENDC_TPL_BOOL_DECL(HEAD_RATIO_ONE, 0, 1),42 ASCENDC_TPL_BOOL_DECL(HEAD_RATIO_ONE, 0, 1),
43- ASCENDC_TPL_BOOL_DECL(BATCH_CONSISTENCY, 0, 1),);43+ ASCENDC_TPL_BOOL_DECL(BATCH_CONSISTENCY, 0, 1),
44+ ASCENDC_TPL_BOOL_DECL(IS_VEC_S2PHYADDR, 0, 1));
44 45 
45// 支持的模板参数组合46// 支持的模板参数组合
46// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法47// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法
@@ -50,7 +51,9 @@ ASCENDC_TPL_SEL(ASCENDC_TPL_ARGS_SEL(
50 ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SMLA_LAYOUT_BSND, SMLA_LAYOUT_TND, SMLA_LAYOUT_PA_BBND),51 ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SMLA_LAYOUT_BSND, SMLA_LAYOUT_TND, SMLA_LAYOUT_PA_BBND),
51 ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, SWA_TEMPLATE, HCA_TEMPLATE, CSA_TEMPLATE,52 ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, SWA_TEMPLATE, HCA_TEMPLATE, CSA_TEMPLATE,
52 ORI_SPARSE_TEMPLATE, ORI_CMP_SPARSE_TEMPLATE),53 ORI_SPARSE_TEMPLATE, ORI_CMP_SPARSE_TEMPLATE),
53- ASCENDC_TPL_BOOL_SEL(SPLIT_G, 0, 1), ASCENDC_TPL_BOOL_SEL(HEAD_RATIO_ONE, 0, 1),54+ ASCENDC_TPL_BOOL_SEL(SPLIT_G, 0, 1),
54- ASCENDC_TPL_BOOL_SEL(BATCH_CONSISTENCY, 0, 1),));55+ ASCENDC_TPL_BOOL_SEL(HEAD_RATIO_ONE, 0, 1),
56+ ASCENDC_TPL_BOOL_SEL(BATCH_CONSISTENCY, 0, 1),
57+ ASCENDC_TPL_BOOL_SEL(IS_VEC_S2PHYADDR, 0, 1)));
55 58 
56#endif // TEMPLATE_TILING_KEY59#endif // TEMPLATE_TILING_KEY
Mattention/sparse_flash_mla/tests/pytest/sparse_flash_mla_paramset.py+72-0
@@ -626,6 +626,78 @@ TEST_PARAMS = {
626 "ori_sparse_indices_mode": ["full"],626 "ori_sparse_indices_mode": ["full"],
627 "cmp_sparse_indices_mode": ["full"],627 "cmp_sparse_indices_mode": ["full"],
628 },628 },
629+ "smla_prefill_pro": {
630+ "testcase_name": [None],
631+ "layout_q": ["TND"],
632+ "layout_kv": ["TND"],
633+ "q_type": [torch.bfloat16],
634+ "ori_kv_type": [torch.bfloat16],
635+ "cmp_kv_type": [torch.bfloat16],
636+ "B": [1],
637+ "S1": [4096],
638+ "S2": [8192],
639+ "T1": [4096],
640+ "T2": [8192],
641+ "T3": [2048],
642+ "N1": [128],
643+ "N2": [1],
644+ "D": [512],
645+ "K": [1024],
646+ "block_num1": [None],
647+ "block_num2": [None],
648+ "block_size1": [128],
649+ "block_size2": [128],
650+ "cu_seqlens_q": [None],
651+ "seqused_ori_kv": [None],
652+ "seqused_cmp_kv": [None],
653+ "cmp_residual_kv": [None],
654+ "softmax_scale": [0.04419417],
655+ "cmp_ratio": [4],
656+ "return_softmax_lse": [False],
657+ "ori_mask_mode": [4],
658+ "cmp_mask_mode": [3],
659+ "ori_win_left": [127],
660+ "ori_win_right": [0],
661+ "template_mode": ["CSA"],
662+ "cmp_kv_topk_mode": ["fullK"],
663+ "cmp_sparse_indices_mode": ["full"],
664+ },
665+ "smla_prefill_flash": {
666+ "testcase_name": [None],
667+ "layout_q": ["TND"],
668+ "layout_kv": ["TND"],
669+ "q_type": [torch.bfloat16],
670+ "ori_kv_type": [torch.bfloat16],
671+ "cmp_kv_type": [torch.bfloat16],
672+ "B": [1],
673+ "S1": [4096],
674+ "S2": [8192],
675+ "T1": [4096],
676+ "T2": [8192],
677+ "T3": [2048],
678+ "N1": [64],
679+ "N2": [1],
680+ "D": [512],
681+ "K": [512],
682+ "block_num1": [None],
683+ "block_num2": [None],
684+ "block_size1": [128],
685+ "block_size2": [128],
686+ "cu_seqlens_q": [None],
687+ "seqused_ori_kv": [None],
688+ "seqused_cmp_kv": [None],
689+ "cmp_residual_kv": [None],
690+ "softmax_scale": [0.04419417],
691+ "cmp_ratio": [4],
692+ "return_softmax_lse": [False],
693+ "ori_mask_mode": [4],
694+ "cmp_mask_mode": [3],
695+ "ori_win_left": [127],
696+ "ori_win_right": [0],
697+ "template_mode": ["CSA"],
698+ "cmp_kv_topk_mode": ["fullK"],
699+ "cmp_sparse_indices_mode": ["full"],
700+ },
629}701}
630 702 
631# 按需选择要启用的测试参数（例如默认启用所有）703# 按需选择要启用的测试参数（例如默认启用所有）
Mattention/sparse_flash_mla/tests/pytest/test_sparse_flash_mla_batch.py+29-12
@@ -23,9 +23,11 @@ import numpy as np
23import os23import os
24 24 
25pt_dir = os.getenv("SMLA_PT_LOAD_PATH", "./data")25pt_dir = os.getenv("SMLA_PT_LOAD_PATH", "./data")
26-result_path = Path(os.getenv("SMLA_RESULT_SAVE_PATH", './result/smla_result.xlsx'))26+result_path = Path(os.getenv("SMLA_RESULT_SAVE_PATH", "./result/smla_result.xlsx"))
27batch_test_mode = int(os.environ.get("SMLA_BATCH_TEST_MODE", 0))27batch_test_mode = int(os.environ.get("SMLA_BATCH_TEST_MODE", 0))
28-excel_path = os.environ.get("SMLA_EXCEL_PATH", os.path.join(os.path.dirname(__file__), "excel", "example.xlsx"))28+excel_path = os.environ.get(
29+ "SMLA_EXCEL_PATH", os.path.join(os.path.dirname(__file__), "excel", "example.xlsx")
30+)
29excel_sheet = os.environ.get("SMLA_EXCEL_SHEET", "CSA")31excel_sheet = os.environ.get("SMLA_EXCEL_SHEET", "CSA")
30 32 
31_single_case_path = os.environ.get("QSAS_TESTCASE_PATH", "").strip()33_single_case_path = os.environ.get("QSAS_TESTCASE_PATH", "").strip()
@@ -33,19 +35,27 @@ _single_case_path = os.environ.get("QSAS_TESTCASE_PATH", "").strip()
33locals()["testcase_files"] = []35locals()["testcase_files"] = []
34if _single_case_path:36if _single_case_path:
35 if not os.path.isfile(_single_case_path):37 if not os.path.isfile(_single_case_path):
36- print(f"错误: 环境变量 QSAS_TESTCASE_PATH 指定的用例文件不存在: {_single_case_path}")38+ print(
39+ f"错误: 环境变量 QSAS_TESTCASE_PATH 指定的用例文件不存在: {_single_case_path}"
40+ )
37 else:41 else:
38 print(f"单用例隔离模式, 仅执行: {_single_case_path}")42 print(f"单用例隔离模式, 仅执行: {_single_case_path}")
39 locals()["testcase_files"].append(_single_case_path)43 locals()["testcase_files"].append(_single_case_path)
40elif os.path.isdir(pt_dir):44elif os.path.isdir(pt_dir):
41- pt_files = [f for f in os.listdir(pt_dir) if f.endswith('.pt')]45+ pt_files = [f for f in os.listdir(pt_dir) if f.endswith(".pt")]
42 if not pt_files:46 if not pt_files:
43 print(f"错误: 目录中没有找到.pt文件: {pt_dir}")47 print(f"错误: 目录中没有找到.pt文件: {pt_dir}")
44 elif batch_test_mode == 1:48 elif batch_test_mode == 1:
45 df = pd.read_excel(excel_path, sheet_name=excel_sheet)49 df = pd.read_excel(excel_path, sheet_name=excel_sheet)
46- target_names = [str(name) for name in df['testcase_name'].dropna().tolist() if str(name) != 'None']50+ target_names = [
51+ str(name)
52+ for name in df["testcase_name"].dropna().tolist()
53+ if str(name) != "None"
54+ ]
47 if not target_names:55 if not target_names:
48- print(f"错误: 表格中没有有效的testcase_name: {excel_path} sheet: {excel_sheet}")56+ print(
57+ f"错误: 表格中没有有效的testcase_name: {excel_path} sheet: {excel_sheet}"
58+ )
49 else:59 else:
50 print(f"从表格[{excel_sheet}]中读取到 {len(target_names)} 个目标用例名")60 print(f"从表格[{excel_sheet}]中读取到 {len(target_names)} 个目标用例名")
51 for target_name in target_names:61 for target_name in target_names:
@@ -68,6 +78,7 @@ else:
68 78 
69print("files:", locals()["testcase_files"])79print("files:", locals()["testcase_files"])
70 80 
81+ 
71@pytest.mark.ci82@pytest.mark.ci
72@pytest.mark.parametrize("testcase_files", locals()["testcase_files"])83@pytest.mark.parametrize("testcase_files", locals()["testcase_files"])
73def test_sparse_flash_mla(testcase_files):84def test_sparse_flash_mla(testcase_files):
@@ -78,22 +89,26 @@ def test_sparse_flash_mla(testcase_files):
78 try:89 try:
79 npu_result, softmax_lse = sparse_flash_mla_process.call_npu(test_data)90 npu_result, softmax_lse = sparse_flash_mla_process.call_npu(test_data)
80 except Exception as e:91 except Exception as e:
81- utils.save_result('Exception', 0, test_data['params'], result_path)92+ utils.save_result("Exception", 0, test_data["params"], result_path)
82 pytest.fail(f"NPU执行异常: {e}")93 pytest.fail(f"NPU执行异常: {e}")
83 94 
84 global_failed = False95 global_failed = False
85 fulfill_percent = 096 fulfill_percent = 0
86 if npu_result is not None:97 if npu_result is not None:
87- result, fulfill_percent = result_compare_method.check_result(test_data['cpu_output'], npu_result)98+ result, fulfill_percent = result_compare_method.check_result(
99+ test_data["cpu_output"], npu_result
100+ )
88 if result == "Failed":101 if result == "Failed":
89 global_failed = True102 global_failed = True
90 else:103 else:
91 global_failed = True104 global_failed = True
92 fulfill_percent = 0105 fulfill_percent = 0
93 106 
94- if test_data['params'].get('return_softmax_lse'):107+ if test_data["params"].get("return_softmax_lse"):
95 print("return_softmax_lse is true!!!")108 print("return_softmax_lse is true!!!")
96- lse_result, lse_percent = result_compare_method.check_result(test_data['softmax_lse'], softmax_lse)109+ lse_result, lse_percent = result_compare_method.check_result(
110+ test_data["softmax_lse"], softmax_lse
111+ )
97 if lse_result == "Failed":112 if lse_result == "Failed":
98 global_failed = True113 global_failed = True
99 fulfill_percent = min(fulfill_percent, lse_percent)114 fulfill_percent = min(fulfill_percent, lse_percent)
@@ -103,6 +118,8 @@ def test_sparse_flash_mla(testcase_files):
103 else:118 else:
104 final_result = "Passed"119 final_result = "Passed"
105 120 
106- utils.save_result(final_result, fulfill_percent, test_data['params'], result_path)121+ utils.save_result(final_result, fulfill_percent, test_data["params"], result_path)
107 if final_result not in ("Passed", "passed", "Pass", "pass"):122 if final_result not in ("Passed", "passed", "Pass", "pass"):
108- pytest.fail(f"用例结果校验失败: result={final_result}, fulfill_percent={fulfill_percent}")123+ pytest.fail(
124+ f"用例结果校验失败: result={final_result}, fulfill_percent={fulfill_percent}"
125+ )
Mattention/sparse_flash_mla/tests/pytest/utils.py+8-8
@@ -217,24 +217,24 @@ def save_result(
217 # 检查文件是否存在217 # 检查文件是否存在
218 result_path.parent.mkdir(parents=True, exist_ok=True)218 result_path.parent.mkdir(parents=True, exist_ok=True)
219 if result_path.exists():219 if result_path.exists():
220- # 读取现有数据
221 df = pd.read_excel(result_path)220 df = pd.read_excel(result_path)
222- 
223- # 检查列名是否一致
224 if set(df.columns) != set(row_data.keys()):221 if set(df.columns) != set(row_data.keys()):
225- print("警告：变量名与Excel列名不匹配！")222+ print("信息：列名不一致，将自动对齐（缺失填NaN，新增列自动扩展）")
226 print(f"Excel列名: {list(df.columns)}")223 print(f"Excel列名: {list(df.columns)}")
227 print(f"变量名: {list(row_data.keys())}")224 print(f"变量名: {list(row_data.keys())}")
228- print("请检查变量名或Excel文件")
229- return False
230- 
231 # 追加新行225 # 追加新行
232 new_df = pd.DataFrame([row_data])226 new_df = pd.DataFrame([row_data])
233- df = pd.concat([df, new_df], ignore_index=True)227+ df = pd.concat([df, new_df], ignore_index=True, sort=False)
234 else:228 else:
235 # 文件不存在，创建新的DataFrame229 # 文件不存在，创建新的DataFrame
236 df = pd.DataFrame([row_data])230 df = pd.DataFrame([row_data])
237 231 
232+ # ---------- 新增：固定最后两列 ----------
233+ # 取出除 result 和 fulfill_percent 外的所有列（保持原有顺序）
234+ other_cols = [col for col in df.columns if col not in ("result", "fulfill_percent")]
235+ # 重新排列：其他列在前，最后两列固定
236+ df = df[other_cols + ["result", "fulfill_percent"]]
237+ # -------------------------------------
238 # 保存到Excel238 # 保存到Excel
239 df.to_excel(result_path, index=False)239 df.to_excel(result_path, index=False)
240 240