Star1294
2968
代码
代码
Issues158
Pull Requests525
Actions
项目讨论
Wiki
分析
Star1294
2968
已合并
feat(ffn_worker_batching/arch35): 下一代完整重写,解耦优化 #9960
zl_hw创建于 8月12日
feat(ffn_worker_batching/arch35): 下一代完整重写,解耦优化 #9960
已合并
zl_hw创建于 8月12日
从zl_hw/ops-transformerfeat/arch35-ffn-a5-rewrite合入到cann/ops-transformermaster
共 13 个文件变更+2241-2090
Mffn/ffn_worker_batching/op_host/ffn_worker_batching_tiling_arch35.cpp+210-37
@@ -12,7 +12,7 @@
12 * \file ffn_worker_batching_tiling_arch35.cpp12 * \file ffn_worker_batching_tiling_arch35.cpp
13 * \brief FfnWorkerBatching arch35 (Ascend950 / DAV_3510) Regbase tiling（1000 档 + IsRegbaseSocVersion 守卫）。13 * \brief FfnWorkerBatching arch35 (Ascend950 / DAV_3510) Regbase tiling（1000 档 + IsRegbaseSocVersion 守卫）。
14 * UB 容量/核数运行时经 GetCoreMemSize/GetCoreNumAiv 取值，禁写死 arch 常量。14 * UB 容量/核数运行时经 GetCoreMemSize/GetCoreNumAiv 取值，禁写死 arch 常量。
15- * 切分阈值/workspace 公式沿用 A2（ffn_worker_batching_tiling.cpp），输入换成 arch35 运行时值自动重算。15+ * 切分阈值与 workspace 布局按 A5 四相位（prepare/sort/gather/group_listing）自行推导。
16 * TilingData 采用 host/kernel 共用平铺 struct（GetTilingData<FfnWorkerBatchingArch35TilingData>() 直写）。16 * TilingData 采用 host/kernel 共用平铺 struct（GetTilingData<FfnWorkerBatchingArch35TilingData>() 直写）。
17 */17 */
18#include "ffn_worker_batching_tiling.h"18#include "ffn_worker_batching_tiling.h"
@@ -46,17 +46,37 @@ constexpr int64_t TILING_KEY_RECV = 101;
46 46 
47constexpr int64_t NUM_TWO = 2;47constexpr int64_t NUM_TWO = 2;
48constexpr int64_t NUM_FOUR = 4;48constexpr int64_t NUM_FOUR = 4;
49+constexpr int64_t NUM_EIGHT = 8;
50+constexpr int64_t GL_ROW_BYTES = static_cast<int64_t>(sizeof(int64_t)) * NUM_TWO; // 一行 = [expert_id, tokenNum]
51+constexpr int64_t GL_UB_FRACTION = 16; // group_list 拼装区取 UB 的 1/16
52+ 
53+// 数据块字节数，与 kernel 侧 AscendC::ONE_BLK_SIZE 同值（host 侧无该符号，故此处按同值定义）。
54+constexpr int64_t ONE_BLK_BYTES = 32;
55+// MrgSort 单轮归并路数：由 MrgSortSrcList 的 4 个入参与 validBit 的 4 个有效位决定。
56+constexpr int64_t MRG_LIST_NUM = 4;
57+// 被 mask 的 token 由上游置为不小于该值的大数，排序前据此压缩剔除。
58+// 与 kernel 侧判据同源（见 op_kernel/ffn_wb_sort_base.h 的 expertStart_ 及算子文档 mask 约定）。
59+constexpr int64_t EXPERT_ID_MASK_START = 1000000;
60+ 
61+// region proposal 对：fp32 键 + uint32 索引，占 SORT_PAIR_FLOATS 个 float。
62+constexpr int64_t SORT_PAIR_FLOATS =
63+ static_cast<int64_t>(sizeof(float) + sizeof(uint32_t)) / static_cast<int64_t>(sizeof(float));
64+// 段内排序时每元素在 UB 的驻留字节：id + 原下标（各 int32）、比较掩码，
65+// 以及 Concat/Sort 要求互不重叠的三块 proposal 对区（concat 结果 / 临时区 / 排序结果）。
66+constexpr int64_t SORT_PAIR_REGIONS = 3;
67+constexpr int64_t SORT_UB_BYTES_PER_ELEM = static_cast<int64_t>(sizeof(int32_t)) * NUM_TWO +
68+ SORT_PAIR_FLOATS * static_cast<int64_t>(sizeof(float)) * SORT_PAIR_REGIONS +
69+ static_cast<int64_t>(sizeof(uint32_t));
70+// 拆包时每元素在 UB 的驻留字节：proposal 对 + 拆出的 id 与 idx。
71+constexpr int64_t EXTRACT_UB_BYTES_PER_ELEM =
72+ SORT_PAIR_FLOATS * static_cast<int64_t>(sizeof(float)) + static_cast<int64_t>(sizeof(int32_t)) * NUM_TWO;
49constexpr int64_t EXPERT_IDX_MAX = 8192;73constexpr int64_t EXPERT_IDX_MAX = 8192;
50constexpr int64_t MAX_SESSION_NUM = 1024;74constexpr int64_t MAX_SESSION_NUM = 1024;
51constexpr int64_t MAX_K_NUM = 64;75constexpr int64_t MAX_K_NUM = 64;
52-constexpr int64_t TH_RECV_CORE_NUM = 32;
53-constexpr int64_t TH_RECV_MIN_ROWS_PER_CORE = 64;
54 76 
55-// arch35 系统预留 UB：GetCoreMemSize(UB) 返回平台标称 UB（本机 248KB），但 vector core77+// 本算子的 SIMT(asc_vf_call) 与 VF 计算需要 UB 系统预留区。预留量不自行相减，
56-// 实际可用比标称少 32KB 系统预留（穿刺2 同进程实测 kernel 侧 UBUF_PER_VECTOR_CORE=216KB）。78+// 而是通过平台接口 ReserveLocalMemory 声明——之后 GetCoreMemSize(UB) 返回的即为可用值。
57-// sort 多核每 loop 4-buffer footprint 按标称 248KB 派生 sortLoopMaxElement 会超物理 UB 溢出，79+// ReservedSize 是平台定义的枚举（8K/16K/32K），选择依据是本算子用到 SIMT,取最大档。
58-// 故派生前扣减系统预留，与同仓 arch35 算子（mhc_pre_sinkhorn_backward）一致。非写死 UB 容量。
59-constexpr uint64_t UB_SYS_RESERVED_SIZE = 32 * 1024;
60} // namespace80} // namespace
61 81 
62class FfnWorkerBatchingTilingArch35 : public Ops::Transformer::OpTiling::TilingBaseClass {82class FfnWorkerBatchingTilingArch35 : public Ops::Transformer::OpTiling::TilingBaseClass {
@@ -83,6 +103,12 @@ protected:
83private:103private:
84 ge::graphStatus CheckInputParam();104 ge::graphStatus CheckInputParam();
85 ge::graphStatus GetAttrsInfo();105 ge::graphStatus GetAttrsInfo();
106+ ge::graphStatus ParseMaxOutShape(const gert::RuntimeAttrs *attrs);
107+ ge::graphStatus ParseOptionalAttrs(const gert::RuntimeAttrs *attrs, int64_t expertNum);
108+ void SplitPrepare();
109+ void SplitSortAndMerge();
110+ void SplitGroupList();
111+ void LayoutWorkspace();
86 112 
87 FfnWorkerBatchingArch35TilingData *tilingDataPtr_ = nullptr;113 FfnWorkerBatchingArch35TilingData *tilingDataPtr_ = nullptr;
88 int64_t A_ = 0;114 int64_t A_ = 0;
@@ -96,6 +122,23 @@ private:
96 int64_t layerNum_ = 0;122 int64_t layerNum_ = 0;
97 int64_t aivNum_ = 0;123 int64_t aivNum_ = 0;
98 int64_t coreNum_ = 0;124 int64_t coreNum_ = 0;
125+ int64_t flatElements_ = 0;
126+ int64_t preparePerLoopRows_ = 0;
127+ int64_t sortSegNum_ = 0;
128+ int64_t sortPerSegElements_ = 0;
129+ int64_t sortLenPerSeg_ = 0;
130+ int64_t mergeRounds_ = 0;
131+ int64_t mergeOneLoopElements_ = 0;
132+ int64_t extractPerLoopElements_ = 0;
133+ int64_t glRowsPerLoop_ = 0;
134+ int64_t bskAlign_ = 0; // BS*K 按数据块对齐后的元素数,DoOpTiling 算好后各切分函数共用
135+ int64_t wsFlatIds_ = 0;
136+ int64_t wsPairA_ = 0;
137+ int64_t wsPairB_ = 0;
138+ int64_t wsSegCnt_ = 0;
139+ int64_t wsSortedIds_ = 0;
140+ int64_t wsGatherIdx_ = 0;
141+ int64_t userWorkspaceWords_ = 0;
99 uint64_t ubSize_ = 0;142 uint64_t ubSize_ = 0;
100 uint32_t sysWorkspaceSize_ = 0;143 uint32_t sysWorkspaceSize_ = 0;
101};144};
@@ -110,13 +153,10 @@ ge::graphStatus FfnWorkerBatchingTilingArch35::GetPlatformInfo()
110 aivNum_ = static_cast<int64_t>(ascendcPlatform.GetCoreNumAiv());153 aivNum_ = static_cast<int64_t>(ascendcPlatform.GetCoreNumAiv());
111 OP_CHECK_IF(aivNum_ == 0, OP_LOGE(context_->GetNodeName(), "Get aivNum failed."), return ge::GRAPH_FAILED);154 OP_CHECK_IF(aivNum_ == 0, OP_LOGE(context_->GetNodeName(), "Get aivNum failed."), return ge::GRAPH_FAILED);
112 155 
156+ // 先声明预留，再取容量：平台在 GetCoreMemSize 中已扣除本次预留，得到 vector core 真实可用 UB。
157+ ascendcPlatform.ReserveLocalMemory(platform_ascendc::ReservedSize::RESERVED_SIZE_32K);
113 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize_);158 ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize_);
114- OP_CHECK_IF(159+ OP_CHECK_IF(ubSize_ == 0, OP_LOGE(context_->GetNodeName(), "Get ubSize failed: 0."), return ge::GRAPH_FAILED);
115- ubSize_ <= UB_SYS_RESERVED_SIZE,
116- OP_LOGE(context_->GetNodeName(), "Get ubSize failed: %lu <= sys reserved %lu.", ubSize_, UB_SYS_RESERVED_SIZE),
117- return ge::GRAPH_FAILED);
118- // 扣减系统预留，得 vector core 实际可用 UB（与 kernel 侧真实口径对齐，防多核 sort footprint 溢出）。
119- ubSize_ -= UB_SYS_RESERVED_SIZE;
120 160 
121 sysWorkspaceSize_ = ascendcPlatform.GetLibApiWorkSpaceSize();161 sysWorkspaceSize_ = ascendcPlatform.GetLibApiWorkSpaceSize();
122 return ge::GRAPH_SUCCESS;162 return ge::GRAPH_SUCCESS;
@@ -152,6 +192,18 @@ ge::graphStatus FfnWorkerBatchingTilingArch35::GetAttrsInfo()
152 return ge::GRAPH_FAILED);192 return ge::GRAPH_FAILED);
153 expertNum_ = *expertNumPtr;193 expertNum_ = *expertNumPtr;
154 194 
195+ if (ParseMaxOutShape(attrs) != ge::GRAPH_SUCCESS) {
196+ return ge::GRAPH_FAILED;
197+ }
198+ if (ParseOptionalAttrs(attrs, *expertNumPtr) != ge::GRAPH_SUCCESS) {
199+ return ge::GRAPH_FAILED;
200+ }
201+ return ge::GRAPH_SUCCESS;
202+}
203+ 
204+// max_out_shape 是 [A, BS, K, H] 四元组:既定形状上界,也是输出 shape 的静态推导依据。
205+ge::graphStatus FfnWorkerBatchingTilingArch35::ParseMaxOutShape(const gert::RuntimeAttrs *attrs)
206+{
155 const gert::ContinuousVector *maxOutShapePtr = attrs->GetAttrPointer<gert::ContinuousVector>(MAX_OUT_SHAPE_ATTR);207 const gert::ContinuousVector *maxOutShapePtr = attrs->GetAttrPointer<gert::ContinuousVector>(MAX_OUT_SHAPE_ATTR);
156 OP_CHECK_NULL_WITH_CONTEXT(context_, maxOutShapePtr);208 OP_CHECK_NULL_WITH_CONTEXT(context_, maxOutShapePtr);
157 OP_CHECK_IF(maxOutShapePtr->GetSize() != static_cast<size_t>(NUM_FOUR),209 OP_CHECK_IF(maxOutShapePtr->GetSize() != static_cast<size_t>(NUM_FOUR),
@@ -176,6 +228,12 @@ ge::graphStatus FfnWorkerBatchingTilingArch35::GetAttrsInfo()
176 228 
177 Y_ = A_ * BS_ * K_;229 Y_ = A_ * BS_ * K_;
178 230 
231+ return ge::GRAPH_SUCCESS;
232+}
233+ 
234+// 可选属性:缺省时保留成员初值,给出即逐个校验取值域。
235+ge::graphStatus FfnWorkerBatchingTilingArch35::ParseOptionalAttrs(const gert::RuntimeAttrs *attrs, int64_t expertNum)
236+{
179 const int64_t *tokenDtype = attrs->GetAttrPointer<int64_t>(TOKEN_DTYPE_ATTR);237 const int64_t *tokenDtype = attrs->GetAttrPointer<int64_t>(TOKEN_DTYPE_ATTR);
180 if (tokenDtype != nullptr) {238 if (tokenDtype != nullptr) {
181 OP_CHECK_IF((*tokenDtype < 0 || *tokenDtype > NUM_TWO),239 OP_CHECK_IF((*tokenDtype < 0 || *tokenDtype > NUM_TWO),
@@ -195,8 +253,8 @@ ge::graphStatus FfnWorkerBatchingTilingArch35::GetAttrsInfo()
195 const int64_t *layNumPtr = attrs->GetAttrPointer<int64_t>(LAY_NUM_ATTR);253 const int64_t *layNumPtr = attrs->GetAttrPointer<int64_t>(LAY_NUM_ATTR);
196 if (layNumPtr != nullptr) {254 if (layNumPtr != nullptr) {
197 OP_CHECK_IF(255 OP_CHECK_IF(
198- (*layNumPtr < 0 || *layNumPtr > *expertNumPtr),256+ (*layNumPtr < 0 || *layNumPtr > expertNum),
199- OP_LOGE(context_->GetNodeName(), "layer_num:%ld must be in range of [0, %ld]", *layNumPtr, *expertNumPtr),257+ OP_LOGE(context_->GetNodeName(), "layer_num:%ld must be in range of [0, %ld]", *layNumPtr, expertNum),
200 return ge::GRAPH_FAILED);258 return ge::GRAPH_FAILED);
201 layerNum_ = *layNumPtr;259 layerNum_ = *layNumPtr;
202 }260 }
@@ -217,18 +275,11 @@ ge::graphStatus FfnWorkerBatchingTilingArch35::DoOpTiling()
217 tilingDataPtr_ = context_->GetTilingData<FfnWorkerBatchingArch35TilingData>();275 tilingDataPtr_ = context_->GetTilingData<FfnWorkerBatchingArch35TilingData>();
218 OP_CHECK_NULL_WITH_CONTEXT(context_, tilingDataPtr_);276 OP_CHECK_NULL_WITH_CONTEXT(context_, tilingDataPtr_);
219 277 
220- coreNum_ = aivNum_;278+ // 用核数由「工作量能否喂饱一个核」决定，不设固定核数阈值：
221- // RECV 限核：A2 用 flat-32（Y<=200000 全覆盖），在 A5（满核 64 + 4T 带宽）下会把带宽 bound 的 gather 腰斩。279+ // 排序按 ONE_REPEAT_SORT_NUM 为粒度推进，一个核至少要分到一个完整粒度才有意义，
222- // A5 策略「只增不减」：Y 大到每核可分满 TH_RECV_MIN_ROWS_PER_CORE 行（gather 占比高、带宽 bound）时放开满核；280+ // 否则多出来的核只是在 SyncAll 上空耗。故上限取平台 aivNum，实际取二者较小值。
223- // Y 较小时 gather 非瓶颈，保留 A2 的 32 核，避免多核 SyncAll 空耗、严格不劣于原实现。281+ const int64_t coreByWork = (Y_ + ONE_REPEAT_SORT_NUM - 1) / ONE_REPEAT_SORT_NUM;
224- // NORM 主线 needSchedule_=0 不触发。A2 路径（monolithic tiling）逐字不动，零回归。282+ coreNum_ = std::max<int64_t>(1, std::min<int64_t>(aivNum_, coreByWork));
225- if (needSchedule_ == 1 && Y_ < aivNum_ * TH_RECV_MIN_ROWS_PER_CORE) {
226- coreNum_ = std::min(aivNum_, TH_RECV_CORE_NUM);
227- }
228- 
229- // UB 派生切分阈值：ubSize 为运行时值（arch35 自动增大），禁写死 arch 常量。
230- int64_t sortLoopMaxElement = static_cast<int64_t>(ubSize_) / (sizeof(int32_t) * NUM_TWO * NUM_FOUR) /
231- ONE_REPEAT_SORT_NUM * ONE_REPEAT_SORT_NUM;
232 283 
233 tilingDataPtr_->Y = Y_;284 tilingDataPtr_->Y = Y_;
234 tilingDataPtr_->H = H_;285 tilingDataPtr_->H = H_;
@@ -236,11 +287,129 @@ ge::graphStatus FfnWorkerBatchingTilingArch35::DoOpTiling()
236 tilingDataPtr_->expertNum = expertNum_;287 tilingDataPtr_->expertNum = expertNum_;
237 tilingDataPtr_->coreNum = coreNum_;288 tilingDataPtr_->coreNum = coreNum_;
238 tilingDataPtr_->ubSize = static_cast<int64_t>(ubSize_);289 tilingDataPtr_->ubSize = static_cast<int64_t>(ubSize_);
239- tilingDataPtr_->sortLoopMaxElement = sortLoopMaxElement;290+ tilingDataPtr_->expertStart = EXPERT_ID_MASK_START;
240- tilingDataPtr_->sortNumWorkSpace = Y_;291+ 
292+ // 扁平序列长度：RECV 的 expert_id 来自 token_info 的 FfnDataDesc，逐 session 取出后按数据块补齐，
293+ // 故为 A*align(BS*K)；NORM 的 expert_ids_buf 本就连续，长度即 Y。
294+ bskAlign_ = (K_ * BS_ * static_cast<int64_t>(sizeof(int32_t)) + ONE_BLK_BYTES - 1) / ONE_BLK_BYTES * ONE_BLK_BYTES /
295+ static_cast<int64_t>(sizeof(int32_t));
296+ flatElements_ = (needSchedule_ == 1) ? A_ * bskAlign_ : Y_;
297+ tilingDataPtr_->flatElements = flatElements_;
298+ 
299+ SplitPrepare();
300+ SplitSortAndMerge();
301+ SplitGroupList();
302+ LayoutWorkspace();
241 return ge::GRAPH_SUCCESS;303 return ge::GRAPH_SUCCESS;
242}304}
243 305 
306+// phase0:单轮块长由运行时 UB 反推。
307+void FfnWorkerBatchingTilingArch35::SplitPrepare()
308+{
309+ // ---------------- phase0(prepare)的切分 ----------------
310+ // RECV 逐 session 行取 BS*K 个 id,同一轮内 UB 需同时驻留:
311+ // · 本轮 id 区 rows * bskAlign * 4B
312+ // · 握手回写的清零区 rows * ONE_BLK_BYTES(每行一个 32B 块写 flag)
313+ // 按运行时 UB 反推每轮行数,不设固定上限;NORM 是整段直搬,按同一公式给出块内元素数即可。
314+ const int64_t prepBytesPerRow = bskAlign_ * static_cast<int64_t>(sizeof(int32_t)) + ONE_BLK_BYTES;
315+ int64_t prepRows = static_cast<int64_t>(ubSize_) / std::max<int64_t>(1, prepBytesPerRow);
316+ prepRows = std::max<int64_t>(1, std::min<int64_t>(prepRows, A_));
317+ preparePerLoopRows_ = prepRows;
318+ tilingDataPtr_->preparePerLoopRows = preparePerLoopRows_;
319+}
320+ 
321+// phase1~3:段内排序(VBS)、段间归并(VMS)与归并收尾(Extract)的切分,三者共用同一份 UB 预算。
322+void FfnWorkerBatchingTilingArch35::SplitSortAndMerge()
323+{
324+ // ---------------- 段内排序（VBS）的切分 ----------------
325+ // 一段在 UB 中同时驻留：输入 id + 原下标（各 4B）、proposal 对区与排序临时区（各 8B/元素）、
326+ // 比较掩码（4B）。故每元素占用 SORT_UB_BYTES_PER_ELEM 字节，据此反推单段元素数上限。
327+ // ubSize_ 在 GetPlatformInfo 中已扣除系统预留，此处直接使用，勿重复扣减。
328+ const int64_t ubAvail = static_cast<int64_t>(ubSize_);
329+ int64_t segCap = ubAvail / SORT_UB_BYTES_PER_ELEM / ONE_REPEAT_SORT_NUM * ONE_REPEAT_SORT_NUM;
330+ segCap = std::max<int64_t>(ONE_REPEAT_SORT_NUM, segCap);
331+ // 先按核数均分；单段超 UB 容量时增加段数（段由各核 grid-stride 认领，段数可多于核数）。
332+ sortSegNum_ = coreNum_;
333+ sortPerSegElements_ = (flatElements_ + sortSegNum_ - 1) / sortSegNum_;
334+ if (sortPerSegElements_ > segCap) {
335+ sortPerSegElements_ = segCap;
336+ sortSegNum_ = (flatElements_ + sortPerSegElements_ - 1) / sortPerSegElements_;
337+ }
338+ sortPerSegElements_ = std::max<int64_t>(1, sortPerSegElements_);
339+ sortSegNum_ = std::max<int64_t>(1, sortSegNum_);
340+ // 每段 proposal 对区：段长按 Sort32 粒度上取整后，每元素占 SORT_PAIR_FLOATS 个 float。
341+ const int64_t segAlign =
342+ (sortPerSegElements_ + ONE_REPEAT_SORT_NUM - 1) / ONE_REPEAT_SORT_NUM * ONE_REPEAT_SORT_NUM;
343+ sortLenPerSeg_ = segAlign * SORT_PAIR_FLOATS;
344+ 
345+ tilingDataPtr_->sortSegNum = sortSegNum_;
346+ tilingDataPtr_->sortPerSegElements = sortPerSegElements_;
347+ tilingDataPtr_->sortLenPerSeg = sortLenPerSeg_;
348+ 
349+ // ---------------- 段间归并（VMS）的切分 ----------------
350+ // 归并轮数：每轮 MRG_LIST_NUM 路合一，直到剩一路。
351+ mergeRounds_ = 0;
352+ for (int64_t lists = sortSegNum_; lists > 1; lists = (lists + MRG_LIST_NUM - 1) / MRG_LIST_NUM) {
353+ mergeRounds_++;
354+ }
355+ // 单次驻留：MRG_LIST_NUM 路输入 + 同宽的输出，均为 proposal 对（每元素 8B）。
356+ // 预算里必须先扣掉同一 TPipe 上的其它缓冲：各段有效数的读回区，以及每个缓冲按块对齐的余量；
357+ // 否则输入与输出两块相加恰好等于可用 UB，分配越界后输出会压到输入上（表现为归并结果头部被覆盖）。
358+ const int64_t mergeOther = sortSegNum_ * ONE_BLK_BYTES + (MRG_LIST_NUM + NUM_TWO) * ONE_BLK_BYTES;
359+ const int64_t mergeUb = (ubAvail > mergeOther) ? (ubAvail - mergeOther) : ubAvail;
360+ int64_t mergeLoop = mergeUb / (MRG_LIST_NUM * NUM_TWO * SORT_PAIR_FLOATS * static_cast<int64_t>(sizeof(float))) /
361+ ONE_REPEAT_SORT_NUM * ONE_REPEAT_SORT_NUM;
362+ mergeOneLoopElements_ = std::max<int64_t>(ONE_REPEAT_SORT_NUM, mergeLoop);
363+ tilingDataPtr_->mergeRounds = mergeRounds_;
364+ tilingDataPtr_->mergeOneLoopElements = mergeOneLoopElements_;
365+ 
366+ // ---------------- 归并收尾（Extract）的切分 ----------------
367+ // 单次驻留：proposal 对（8B）+ 拆出的 id 与 idx（各 4B）。
368+ int64_t extractLoop = ubAvail / EXTRACT_UB_BYTES_PER_ELEM / ONE_REPEAT_SORT_NUM * ONE_REPEAT_SORT_NUM;
369+ extractPerLoopElements_ = std::max<int64_t>(ONE_REPEAT_SORT_NUM, std::min<int64_t>(extractLoop, Y_));
370+ tilingDataPtr_->extractPerLoopElements = extractPerLoopElements_;
371+}
372+ 
373+// phase4:group_list 写出时每块拼多少行。
374+void FfnWorkerBatchingTilingArch35::SplitGroupList()
375+{
376+ // ---------------- group_list 的切分 ----------------
377+ // 每行 [expert_id, tokenNum] 两个 int64 = GL_ROW_BYTES；拼装区取 UB 的 1/GL_UB_FRACTION，
378+ // 按 2 行对齐（使块起点落在数据块边界），并以 expertNum 封顶。
379+ int64_t rows = static_cast<int64_t>(ubSize_) / GL_UB_FRACTION / GL_ROW_BYTES / NUM_TWO * NUM_TWO;
380+ glRowsPerLoop_ = std::max<int64_t>(NUM_TWO, std::min<int64_t>(rows, expertNum_));
381+ 
382+ tilingDataPtr_->glRowsPerLoop = glRowsPerLoop_;
383+}
384+ 
385+// workspace 段偏移:与 GetWorkspaceSize 的累加顺序严格一致,两处取自同一组成员变量。
386+void FfnWorkerBatchingTilingArch35::LayoutWorkspace()
387+{
388+ // ---------------- workspace 段偏移（以 int32 word 计）----------------
389+ // 布局与 GetWorkspaceSize 的累加顺序严格一致，两处取自同一组成员变量。
390+ int64_t off = MAX_RESERVE_WK_NUM;
391+ wsFlatIds_ = off;
392+ off += flatElements_;
393+ wsPairA_ = off;
394+ off += sortSegNum_ * sortLenPerSeg_;
395+ wsPairB_ = off;
396+ off += sortSegNum_ * sortLenPerSeg_;
397+ wsSegCnt_ = off;
398+ off += sortSegNum_ * (ONE_BLK_BYTES / static_cast<int64_t>(sizeof(int32_t)));
399+ wsSortedIds_ = off;
400+ off += Y_;
401+ wsGatherIdx_ = off;
402+ off += Y_;
403+ userWorkspaceWords_ = off;
404+ 
405+ tilingDataPtr_->wsFlatIds = wsFlatIds_;
406+ tilingDataPtr_->wsPairA = wsPairA_;
407+ tilingDataPtr_->wsPairB = wsPairB_;
408+ tilingDataPtr_->wsSegCnt = wsSegCnt_;
409+ tilingDataPtr_->wsSortedIds = wsSortedIds_;
410+ tilingDataPtr_->wsGatherIdx = wsGatherIdx_;
411+}
412+ 
244ge::graphStatus FfnWorkerBatchingTilingArch35::DoLibApiTiling()413ge::graphStatus FfnWorkerBatchingTilingArch35::DoLibApiTiling()
245{414{
246 return ge::GRAPH_SUCCESS;415 return ge::GRAPH_SUCCESS;
@@ -254,9 +423,9 @@ uint64_t FfnWorkerBatchingTilingArch35::GetTilingKey() const
254 423 
255ge::graphStatus FfnWorkerBatchingTilingArch35::GetWorkspaceSize()424ge::graphStatus FfnWorkerBatchingTilingArch35::GetWorkspaceSize()
256{425{
257- // 与 A2 同公式；sysWorkspaceSize 经 GetLibApiWorkSpaceSize() 取 arch35 平台值。426+ // 用户区总量由 DoOpTiling 逐段累加得到（userWorkspaceWords_），此处不另算一遍：
258- workspaceSize_ = MAX_RESERVE_WK_NUM * sizeof(int32_t) + Y_ * sizeof(int32_t) +427+ // 段偏移与总量出自同一次累加，避免布局在两处各写一遍而悄悄错位。
259- Y_ * sizeof(int32_t) * NUM_TWO * NUM_FOUR + expertNum_ * sizeof(int32_t) + sysWorkspaceSize_;428+ workspaceSize_ = userWorkspaceWords_ * static_cast<int64_t>(sizeof(int32_t)) + sysWorkspaceSize_;
260 return ge::GRAPH_SUCCESS;429 return ge::GRAPH_SUCCESS;
261}430}
262 431 
@@ -270,11 +439,15 @@ ge::graphStatus FfnWorkerBatchingTilingArch35::PostTiling()
270 currentWorkspace[0] = static_cast<size_t>(workspaceSize_);439 currentWorkspace[0] = static_cast<size_t>(workspaceSize_);
271 440 
272 OP_LOGI(context_->GetNodeName(),441 OP_LOGI(context_->GetNodeName(),
273- "arch35 tiling: coreNum:%ld ubSize:%ld Y:%ld H:%ld tokenDtype:%ld expertNum:%ld "442+ "arch35 tiling: coreNum:%ld ubSize:%ld Y:%ld H:%ld tokenDtype:%ld expertNum:%ld flatElements:%ld "
274- "sortLoopMaxElement:%ld sortNumWorkSpace:%ld tilingKey:%lu",443+ "prepRows:%ld "
444+ "sortSegNum:%ld sortPerSeg:%ld sortLenPerSeg:%ld mergeRounds:%ld mergeLoop:%ld extractLoop:%ld "
445+ "glRows:%ld wsWords:%ld tilingKey:%lu",
275 tilingDataPtr_->coreNum, tilingDataPtr_->ubSize, tilingDataPtr_->Y, tilingDataPtr_->H,446 tilingDataPtr_->coreNum, tilingDataPtr_->ubSize, tilingDataPtr_->Y, tilingDataPtr_->H,
276- tilingDataPtr_->tokenDtype, tilingDataPtr_->expertNum, tilingDataPtr_->sortLoopMaxElement,447+ tilingDataPtr_->tokenDtype, tilingDataPtr_->expertNum, tilingDataPtr_->flatElements,
277- tilingDataPtr_->sortNumWorkSpace, GetTilingKey());448+ tilingDataPtr_->preparePerLoopRows, tilingDataPtr_->sortSegNum, tilingDataPtr_->sortPerSegElements,
449+ tilingDataPtr_->sortLenPerSeg, tilingDataPtr_->mergeRounds, tilingDataPtr_->mergeOneLoopElements,
450+ tilingDataPtr_->extractPerLoopElements, tilingDataPtr_->glRowsPerLoop, userWorkspaceWords_, GetTilingKey());
278 return ge::GRAPH_SUCCESS;451 return ge::GRAPH_SUCCESS;
279}452}
280 453 
Affn/ffn_worker_batching/op_kernel/arch35/ffn_wb_a5_context.h+189-0
@@ -0,0 +1,189 @@
1+/**
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
3+ *
4+ * Licensed under the Apache License, Version 2.0 (the "License");
5+ * you may not use this file except in compliance with the License.
6+ * You may obtain a copy of the License at
7+ *
8+ * http://www.apache.org/licenses/LICENSE-2.0
9+ *
10+ * Unless required by applicable law or agreed to in writing, software
11+ * distributed under the License is distributed on an "AS IS" BASIS,
12+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+ * See the License for the specific language governing permissions and
14+ * limitations under the License.
15+ */
16+ 
17+/*!
18+ * \file ffn_wb_a5_context.h
19+ * \brief arch35(A5) 自有基础层:输入契约解析 + 工具函数。A5 不再经桥接头引用 A2 的
20+ * ../ffn_wb_common.h 与 ../ffn_wb_get_schedule_context.h。
21+ *
22+ * 本文件内**不写任何裸常量**:
23+ * · schedule_context 的字段偏移与结构体大小,取自公共契约头 attention_ffn_schedule.h,
24+ * 由编译器 offsetof/sizeof 推出(A2 侧那份 valLocal[528] 之类的手写偏移是待淘汰写法);
25+ * · 数据块字节数用 AscendC 的 ONE_BLK_SIZE,不自定义 32;
26+ * · proposal 对宽度由"fp32 键 + uint32 索引"两个类型宽度相加得到,不写 8;
27+ * · fp32 下界由编译器内建 __FLT_MAX__ 取反得到,不写 -3.4e38;
28+ * · 切分/workspace 段偏移一律由 host tiling 下发,kernel 不持有任何布局常量;
29+ * · A/M/K 的上界由 host 侧 CheckInputParam 统一校验,kernel 不再重复声明限值。
30+ */
31+#ifndef OP_KERNEL_ARCH35_FFN_WB_A5_CONTEXT_H
32+#define OP_KERNEL_ARCH35_FFN_WB_A5_CONTEXT_H
33+#include "kernel_operator.h"
34+#include "attention_ffn_schedule.h" // 与 Attention 侧约定的 schedule_context 权威定义
35+ 
36+namespace FfnWbBatchingArch35 {
37+using namespace AscendC;
38+ 
39+// 契约字段在 schedule_context 中的字节偏移:由权威结构体推出,不手写。
40+#define FFN_WB_CTX_OFFSET(field) static_cast<int32_t>(__builtin_offsetof(aicpu::ScheduleContext, field))
41+// schedule_context 的字节长度同样取自权威结构体(该结构体自带 static_assert 保证为约定值)。
42+constexpr int64_t SCHEDULE_CONTEXT_BYTES = static_cast<int64_t>(sizeof(aicpu::ScheduleContext));
43+ 
44+// 排序中间表示:region proposal 对 = fp32 键 + uint32 索引。
45+constexpr int64_t SORT_PAIR_BYTES = static_cast<int64_t>(sizeof(float) + sizeof(uint32_t));
46+// 降序排序中代表"最小"的填充值:取 fp32 可表示的最小有限值。
47+constexpr float SORT_FILL_VALUE = -__FLT_MAX__;
48+// 以下两个是**硬件指令粒度**,无法由其他量推导:
49+// · Sort32 单次排序元素数:SDK 定义在 impl/basic_api/dav_3510/kernel_operator_proposal_impl.h
50+// (singleSortElementCountArch3510 = 32),该头为内部实现头、禁止直接 include,故此处按同值定义;
51+// · MrgSort 单轮归并路数:由 MrgSortSrcList 的 4 个入参与 validBit 的 4 个有效位决定。
52+// 二者若随架构变化,以 SDK 上述定义为准。
53+constexpr int64_t ONE_REPEAT_SORT_NUM = 32;
54+constexpr int64_t MRG_LIST_NUM = 4;
55+ 
56+// schedule_context 内存放的是设备地址(二级指针),真数据需二次解引用后使用。
57+struct BufferInfo {
58+ uint64_t tokenInfoBuf = 0;
59+ uint64_t tokenDataBuf = 0;
60+ uint64_t sessionIdsBuf = 0;
61+ uint64_t microBatchIdsBuf = 0;
62+ uint64_t expertIdsBuf = 0;
63+};
64+ 
65+struct ScheduleContextInfo {
66+ uint32_t A = 0; // attention session num
67+ uint32_t M = 0; // micro batch num
68+ uint32_t BS = 0; // micro batch size
69+ uint32_t K = 0; // selected expert num(topK+1)
70+ uint32_t HS = 0; // attn_to_ffn_token_size,单位字节
71+ uint32_t H = 0; // hidden size(attr)
72+ uint32_t Y = 0; // A*BS*K
73+ uint64_t curMicroBatchID = 0; // RECV:当前 expert id 已就绪的 micro batch
74+ uint32_t outNum = 0; // NORM:FfnArea 中有效的 session 数
75+ uint32_t tokenDtype = 0; // 0:FP16 1:BF16 2:int8 + dynamic scale 连续排布
76+ uint32_t expertNum = 0;
77+ int64_t coreNum = 0;
78+ int64_t ubSize = 0;
79+ int64_t validGatherIdxLength = 0; // 排序后剔除无效值的有效长度,≤ A*BS*K
80+ int64_t BsKPaddingCount = 0; // RECV:BS*K 按 block 对齐需补的个数
81+ 
82+ BufferInfo bufferPtr;
83+};
84+ 
85+// ---------------- 工具函数 ----------------
86+template <HardEvent event>
87+__aicore__ inline void SetWaitFlag(HardEvent evt)
88+{
89+ event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(evt));
90+ SetFlag<event>(eventId);
91+ WaitFlag<event>(eventId);
92+}
93+ 
94+__aicore__ inline int64_t Ceil(int64_t a, int64_t b)
95+{
96+ return (b == 0) ? 0 : (a + b - 1) / b;
97+}
98+ 
99+// 与 Ceil 同义,保留此名以贴合"向上取整除"的调用点语义。
100+__aicore__ inline int64_t CeilDiv(int64_t a, int64_t b)
101+{
102+ return Ceil(a, b);
103+}
104+ 
105+// 数据块字节数别名:沿用 AscendC 的 ONE_BLK_SIZE,不另立常量。
106+constexpr int64_t BLOCK_BYTES = static_cast<int64_t>(ONE_BLK_SIZE);
107+constexpr int64_t NUM_TWO = 2;
108+constexpr int64_t NUM_THREE = 3;
109+constexpr int64_t NUM_FOUR = 4;
110+ 
111+// 返回按一个数据块(ONE_BLK_SIZE)对齐后的**元素个数**,bytes 为单元素字节数。
112+__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes)
113+{
114+ return (bytes == 0) ? 0 : (elementNum * bytes + ONE_BLK_SIZE - 1) / ONE_BLK_SIZE * ONE_BLK_SIZE / bytes;
115+}
116+ 
117+template <typename T>
118+__aicore__ inline T Min(T a, T b)
119+{
120+ return (a < b) ? a : b;
121+}
122+ 
123+template <typename T>
124+__aicore__ inline T Max(T a, T b)
125+{
126+ return (a > b) ? a : b;
127+}
128+ 
129+// proposal 对表示下,count 个元素占用的 T 元素个数。
130+template <typename T>
131+__aicore__ inline int64_t GetSortLen(int64_t count)
132+{
133+ return count * SORT_PAIR_BYTES / static_cast<int64_t>(sizeof(T));
134+}
135+ 
136+// ---------------- schedule_context 解析 ----------------
137+// isRecv=false(NORM):expert_id 取自 FfnArea.expert_ids_buf,并读 session/micro_batch ids 与 out_num;
138+// isRecv=true (RECV):expert_id 藏在 FfnArea.token_info_buf 的 FfnDataDesc 内,
139+// 另读 polling_index 定位当前 micro batch。
140+template <bool isRecv = false, typename TilingT>
141+__aicore__ inline void ScheduleContextParse(GM_ADDR schedule_context, const TilingT *tilingData,
142+ ScheduleContextInfo &ctx, TPipe *pipe)
143+{
144+ GlobalTensor<int8_t> ctxGm;
145+ TBuf<TPosition::VECIN> buffer;
146+ ctxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int8_t *>(schedule_context), SCHEDULE_CONTEXT_BYTES);
147+ pipe->InitBuffer(buffer, SCHEDULE_CONTEXT_BYTES);
148+ LocalTensor<int8_t> val = buffer.Get<int8_t>();
149+ DataCopy(val, ctxGm, SCHEDULE_CONTEXT_BYTES);
150+ SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
151+ 
152+ ctx.A = val[FFN_WB_CTX_OFFSET(common.session_num)].template ReinterpretCast<uint32_t>().GetValue(0);
153+ ctx.M = val[FFN_WB_CTX_OFFSET(common.micro_batch_num)].template ReinterpretCast<uint32_t>().GetValue(0);
154+ ctx.BS = val[FFN_WB_CTX_OFFSET(common.micro_batch_size)].template ReinterpretCast<uint32_t>().GetValue(0);
155+ ctx.K = val[FFN_WB_CTX_OFFSET(common.selected_expert_num)].template ReinterpretCast<uint32_t>().GetValue(0);
156+ ctx.HS = val[FFN_WB_CTX_OFFSET(common.attn_to_ffn_token_size)].template ReinterpretCast<uint32_t>().GetValue(0);
157+ 
158+ ctx.H = tilingData->H;
159+ ctx.Y = tilingData->Y;
160+ ctx.tokenDtype = tilingData->tokenDtype;
161+ ctx.expertNum = tilingData->expertNum;
162+ ctx.coreNum = tilingData->coreNum;
163+ ctx.ubSize = tilingData->ubSize;
164+ 
165+ ctx.bufferPtr.tokenDataBuf =
166+ val[FFN_WB_CTX_OFFSET(ffn.token_data_buf)].template ReinterpretCast<uint64_t>().GetValue(0);
167+ 
168+ if constexpr (isRecv) {
169+ ctx.bufferPtr.tokenInfoBuf =
170+ val[FFN_WB_CTX_OFFSET(ffn.token_info_buf)].template ReinterpretCast<uint64_t>().GetValue(0);
171+ ctx.curMicroBatchID =
172+ val[FFN_WB_CTX_OFFSET(ffn.polling_index)].template ReinterpretCast<uint64_t>().GetValue(0);
173+ ASSERT_MSG(ctx.curMicroBatchID < ctx.M, "curMicroBatchID:%lu should be less than micro_batch_num:%u",
174+ ctx.curMicroBatchID, ctx.M);
175+ const int64_t bsk = static_cast<int64_t>(ctx.BS) * ctx.K;
176+ ctx.BsKPaddingCount = Align(bsk, sizeof(int32_t)) - bsk;
177+ } else {
178+ ctx.bufferPtr.sessionIdsBuf =
179+ val[FFN_WB_CTX_OFFSET(ffn.session_ids_buf)].template ReinterpretCast<uint64_t>().GetValue(0);
180+ ctx.bufferPtr.microBatchIdsBuf =
181+ val[FFN_WB_CTX_OFFSET(ffn.micro_batch_ids_buf)].template ReinterpretCast<uint64_t>().GetValue(0);
182+ ctx.bufferPtr.expertIdsBuf =
183+ val[FFN_WB_CTX_OFFSET(ffn.expert_ids_buf)].template ReinterpretCast<uint64_t>().GetValue(0);
184+ ctx.outNum = val[FFN_WB_CTX_OFFSET(ffn.out_num)].template ReinterpretCast<uint32_t>().GetValue(0);
185+ }
186+}
187+ 
188+} // namespace FfnWbBatchingArch35
189+#endif // OP_KERNEL_ARCH35_FFN_WB_A5_CONTEXT_H
Affn/ffn_worker_batching/op_kernel/arch35/ffn_wb_a5_gather.h+447-0
@@ -0,0 +1,447 @@
1+/**
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
3+ *
4+ * Licensed under the Apache License, Version 2.0 (the "License");
5+ * you may not use this file except in compliance with the License.
6+ * You may obtain a copy of the License at
7+ *
8+ * http://www.apache.org/licenses/LICENSE-2.0
9+ *
10+ * Unless required by applicable law or agreed to in writing, software
11+ * distributed under the License is distributed on an "AS IS" BASIS,
12+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+ * See the License for the specific language governing permissions and
14+ * limitations under the License.
15+ */
16+ 
17+/*!
18+ * \file ffn_wb_a5_gather.h
19+ * \brief arch35(A5) 按 gather_idx 取行搬运:y / session_ids / micro_batch_ids / token_ids /
20+ * expert_offsets / dynamic_scale。
21+ *
22+ * 逐行取 idx 到标量再算源地址、每行走 DataCopyPad 搬 H 字节 —— 这是 A5 上同类"按索引取行"
23+ * 的惯用形态(同仓 moe_v3_gather_out.h 的 arch35 实现同构):行数据是大块搬运,应交给 MTE;
24+ * 而源地址由数据决定,索引必须落到标量才能喂给搬运指令。
25+ */
26+#ifndef OP_KERNEL_ARCH35_FFN_WB_A5_GATHER_H
27+#define OP_KERNEL_ARCH35_FFN_WB_A5_GATHER_H
28+#include "ffn_wb_a5_context.h"
29+ 
30+namespace FfnWbBatchingArch35 {
31+using namespace AscendC;
32+ 
33+// 索引换算的向量分道:每个中间量独占一道,不复用、不自我覆盖(PER_LOOP_ROWS=128,每道 512B)。
34+constexpr int64_t IDX_LANE_SRC = 0; // gatherIdx(fp32)
35+constexpr int64_t IDX_LANE_QA = 1; // gatherIdx / bskProduct
36+constexpr int64_t IDX_LANE_AF = 2; // float(aIdx)
37+constexpr int64_t IDX_LANE_REM = 3; // gatherIdx - aIdx*bskProduct
38+constexpr int64_t IDX_LANE_QB = 4; // rem / K
39+constexpr int64_t IDX_LANE_BF = 5; // float(bsIdx)
40+constexpr int64_t IDX_LANE_KF = 6; // rem - bsIdx*K
41+constexpr int64_t IDX_LANE_AI = 0; // aIdx(int32,独立缓冲)
42+constexpr int64_t IDX_LANE_BI = 1; // bsIdx
43+constexpr int64_t IDX_LANE_KI = 2; // kIdx
44+constexpr int64_t IDX_F_LANES = 7; // srcF/qaF/aF/remF/qbF/bF/kF
45+constexpr int64_t IDX_I_LANES = 3; // aIdx/bsIdx/kIdx
46+ 
47+// ===================== step 2:gather 搬运(原 A2 ffn_wb_gather_out_all.h) =====================
48+template <bool isScanFlag = false>
49+class FfnWbA5Gather {
50+public:
51+ __aicore__ inline FfnWbA5Gather()
52+ {
53+ }
54+ __aicore__ inline void Init(GM_ADDR expertid_idx, GM_ADDR y, GM_ADDR session_ids, GM_ADDR micro_batch_ids,
55+ GM_ADDR token_ids, GM_ADDR expert_offsets, GM_ADDR dynamic_scale,
56+ const ScheduleContextInfo *contextInfo, TPipe *pipe, uint32_t usedCoreNum)
57+ {
58+ contextInfo_ = contextInfo;
59+ curMicroBatchID = contextInfo_->curMicroBatchID;
60+ BsKPaddingCount = contextInfo_->BsKPaddingCount;
61+ int64_t useCore = contextInfo_->coreNum - usedCoreNum;
62+ 
63+ tokenDtypeSize_ = (contextInfo_->tokenDtype == NUM_TWO) ? sizeof(int8_t) : sizeof(half); // attr 2: int8
64+ 
65+ sessionNumBlockAlign_ = Align(contextInfo_->A, sizeof(int32_t));
66+ int64_t validGatherIdxLength = contextInfo_->validGatherIdxLength;
67+ 
68+ // ⚠️ 逐项扣除本类**所有**会 InitBuffer 的缓冲。少扣一项会让 maxBlockSize_ 算大,
69+ // inQueueX_ 随即申请越界、分配失败,输出成为垃圾。新增缓冲必须同步加进这个式子。
70+ int64_t ubAvailable =
71+ contextInfo->ubSize -
72+ (BUFFER_NUM * PER_LOOP_ROWS * sizeof(int32_t) * VAR_NUM + BUFFER_NUM * PER_LOOP_ROWS * sizeof(int32_t) +
73+ sessionNumBlockAlign_ * sizeof(int32_t) * BUFFER_NUM + PER_LOOP_ROWS * BLOCK_SIZE * BUFFER_NUM +
74+ PER_LOOP_ROWS * (sizeof(float) * IDX_F_LANES + sizeof(int32_t) * IDX_I_LANES));
75+ 
76+ int64_t maxTokenSize = ubAvailable / BUFFER_NUM;
77+ maxBlockSize_ = maxTokenSize - (contextInfo_->tokenDtype == TOKEN_KIND_TWO ? BLOCK_BYTES : 0);
78+ maxBlockSize_ = (maxBlockSize_ / BLOCK_BYTES * BLOCK_BYTES) / tokenDtypeSize_;
79+ hBlocks_ = (contextInfo_->H + maxBlockSize_ - 1) / maxBlockSize_;
80+ lastHBlockSize_ = contextInfo_->H - (hBlocks_ - 1) * maxBlockSize_;
81+ 
82+ int64_t blockIdx = GetBlockIdx();
83+ int64_t perCoreRows = CeilDiv(validGatherIdxLength, useCore);
84+ needCoreNum_ = perCoreRows == 0 ? 0 : CeilDiv(validGatherIdxLength, perCoreRows);
85+ int64_t lastCoreRows = validGatherIdxLength - perCoreRows * (needCoreNum_ - 1);
86+ 
87+ if (blockIdx == needCoreNum_ - 1) {
88+ lastLoopRows_ = lastCoreRows - (CeilDiv(lastCoreRows, PER_LOOP_ROWS) - 1) * PER_LOOP_ROWS;
89+ rowLoops_ = (lastCoreRows + PER_LOOP_ROWS - 1) / PER_LOOP_ROWS;
90+ } else {
91+ lastLoopRows_ = perCoreRows - (CeilDiv(perCoreRows, PER_LOOP_ROWS) - 1) * PER_LOOP_ROWS;
92+ rowLoops_ = (perCoreRows + PER_LOOP_ROWS - 1) / PER_LOOP_ROWS;
93+ }
94+ uint64_t SplitY = perCoreRows * contextInfo_->H * tokenDtypeSize_;
95+ 
96+ GM_ADDR tokenDataBufAddr = reinterpret_cast<GM_ADDR>(contextInfo_->bufferPtr.tokenDataBuf);
97+ GM_ADDR sessionIdsBufAddr = reinterpret_cast<GM_ADDR>(contextInfo_->bufferPtr.sessionIdsBuf);
98+ GM_ADDR microBatchIdsBufAddr = reinterpret_cast<GM_ADDR>(contextInfo_->bufferPtr.microBatchIdsBuf);
99+ 
100+ tokenDataBufGm_.SetGlobalBuffer((__gm__ int8_t *)tokenDataBufAddr);
101+ 
102+ // 排序的后的 对应gather_index
103+ expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertid_idx + blockIdx * perCoreRows);
104+ 
105+ sessionIdsInGm_.SetGlobalBuffer((__gm__ int32_t *)sessionIdsBufAddr, contextInfo_->A);
106+ microBatchIdsInGm_.SetGlobalBuffer((__gm__ int32_t *)microBatchIdsBufAddr, contextInfo_->A);
107+ 
108+ // 输出空间
109+ yOutGm_.SetGlobalBuffer((__gm__ int8_t *)y + blockIdx * SplitY);
110+ 
111+ sessionIdsOutGm_.SetGlobalBuffer((__gm__ int32_t *)session_ids + blockIdx * perCoreRows);
112+ microBatchIdsOutGm_.SetGlobalBuffer((__gm__ int32_t *)micro_batch_ids + blockIdx * perCoreRows);
113+ tokenIdsOutGm_.SetGlobalBuffer((__gm__ int32_t *)token_ids + blockIdx * perCoreRows);
114+ expertOffsetsOutGm_.SetGlobalBuffer((__gm__ int32_t *)expert_offsets + blockIdx * perCoreRows);
115+ if (contextInfo_->tokenDtype == TOKEN_KIND_TWO) {
116+ dynamicScaleOutGm_.SetGlobalBuffer((__gm__ float *)dynamic_scale + blockIdx * perCoreRows);
117+ }
118+ 
119+ InitBuffers(pipe);
120+ }
121+ 
122+ // 各缓冲的开辟:块大小已由上面按可用 UB 逐项扣除算好,这里只做分配。
123+ __aicore__ inline void InitBuffers(TPipe *pipe)
124+ {
125+ int64_t blockBufferSize =
126+ maxBlockSize_ * tokenDtypeSize_ + (contextInfo_->tokenDtype == TOKEN_KIND_TWO ? BLOCK_BYTES : 0);
127+ pipe->InitBuffer(inQueueX_, BUFFER_NUM, blockBufferSize);
128+ 
129+ // PER_LOOP_ROWS 为长度 包含 额外5 + 1个输出;
130+ pipe->InitBuffer(outQueALL_, BUFFER_NUM,
131+ PER_LOOP_ROWS * sizeof(int32_t) * VAR_NUM + PER_LOOP_ROWS * BLOCK_SIZE);
132+ 
133+ // 将gm gather_idx 一段长度 放到 UB 空间的
134+ pipe->InitBuffer(expertIdxQue_, BUFFER_NUM, PER_LOOP_ROWS * sizeof(int32_t));
135+ // 索引换算的向量分道(已计入上面的 ubAvailable)
136+ pipe->InitBuffer(idxCalcBuf_, PER_LOOP_ROWS * sizeof(float) * IDX_F_LANES);
137+ pipe->InitBuffer(idxIntBuf_, PER_LOOP_ROWS * sizeof(int32_t) * IDX_I_LANES);
138+ 
139+ // 将gm buf 放到 UB 空间的
140+ if constexpr (isScanFlag == false) {
141+ pipe->InitBuffer(tmpBuffer_, sessionNumBlockAlign_ * sizeof(int32_t) * BUFFER_NUM);
142+ }
143+ }
144+ 
145+ __aicore__ inline void CopyInIds()
146+ {
147+ DataCopyExtParams copyParams1{1, static_cast<uint32_t>(contextInfo_->A * sizeof(int32_t)), 0, 0, 0};
148+ DataCopyPadExtParams<int32_t> padParams1{false, 0, 0, 0};
149+ sessionIdsLocal_ = tmpBuffer_.Get<int32_t>();
150+ microBatchIdsLocal_ = sessionIdsLocal_[sessionNumBlockAlign_];
151+ 
152+ DataCopyPad(sessionIdsLocal_, sessionIdsInGm_, copyParams1, padParams1);
153+ DataCopyPad(microBatchIdsLocal_, microBatchIdsInGm_, copyParams1, padParams1);
154+ }
155+ 
156+ __aicore__ inline void CopyInExpertIdx(int32_t expertIdxOffset, int32_t curRows)
157+ {
158+ LocalTensor<int32_t> expertIdxLocal = expertIdxQue_.AllocTensor<int32_t>();
159+ DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(curRows * sizeof(int32_t)), 0, 0,
160+ 0};
161+ DataCopyPadExtParams<int32_t> padParams{false, 0, 0, 0};
162+ DataCopyPad(expertIdxLocal, expertIdxGm_[expertIdxOffset], copyParams, padParams);
163+ expertIdxQue_.EnQue(expertIdxLocal);
164+ }
165+ 
166+ __aicore__ inline void Process()
167+ {
168+ if (GetBlockIdx() >= needCoreNum_) {
169+ return;
170+ }
171+ 
172+ int64_t bskProduct = contextInfo_->BS * contextInfo_->K;
173+ if constexpr (isScanFlag == false) {
174+ CopyInIds();
175+ } else {
176+ bskProduct = contextInfo_->BS * contextInfo_->K + BsKPaddingCount;
177+ }
178+ 
179+ int64_t curLoopElements = PER_LOOP_ROWS;
180+ int64_t strideSession = contextInfo_->M * contextInfo_->BS * contextInfo_->K * contextInfo_->HS;
181+ int64_t strideMicroBatch = contextInfo_->BS * contextInfo_->K * contextInfo_->HS;
182+ int64_t strideBs = contextInfo_->K * contextInfo_->HS;
183+ int64_t strideK = contextInfo_->HS;
184+ 
185+ for (int64_t i = 0; i < rowLoops_; i++) {
186+ int64_t currentOuterStart = i * PER_LOOP_ROWS;
187+ if (i == rowLoops_ - 1) {
188+ curLoopElements = lastLoopRows_;
189+ }
190+ 
191+ CopyInExpertIdx(currentOuterStart, curLoopElements);
192+ LocalTensor<int32_t> expertIdxLocal = expertIdxQue_.DeQue<int32_t>();
193+ LocalTensor<int32_t> outAllLocal = outQueALL_.AllocTensor<int32_t>();
194+ SetWaitFlag<HardEvent::MTE2_V>(HardEvent::MTE2_V);
195+ // 下标换算整块向量化,结果落在 aIdxAll / bsIdxAll / kIdxAll 三路分道
196+ DecomposeIndices(expertIdxLocal, bskProduct, curLoopElements);
197+ LocalTensor<int32_t> ints = idxIntBuf_.Get<int32_t>();
198+ LocalTensor<int32_t> aIdxAll = ints[PER_LOOP_ROWS * IDX_LANE_AI];
199+ LocalTensor<int32_t> bsIdxAll = ints[PER_LOOP_ROWS * IDX_LANE_BI];
200+ LocalTensor<int32_t> kIdxAll = ints[PER_LOOP_ROWS * IDX_LANE_KI];
201+ 
202+ EmitTokens(outAllLocal, curLoopElements, currentOuterStart, strideSession, strideMicroBatch, strideBs,
203+ strideK, aIdxAll, bsIdxAll, kIdxAll);
204+ 
205+ outQueALL_.EnQue(outAllLocal);
206+ CopyAllLocalOut(currentOuterStart, curLoopElements);
207+ expertIdxQue_.FreeTensor(expertIdxLocal);
208+ }
209+ }
210+ 
211+ // 逐 token:从三路分道取出 (a, bs, k),算出源地址搬 H 字节,并填四个 id 输出。
212+ __aicore__ inline void EmitTokens(const LocalTensor<int32_t> &outAllLocal, int64_t curLoopElements,
213+ int64_t currentOuterStart, int64_t strideSession, int64_t strideMicroBatch,
214+ int64_t strideBs, int64_t strideK, const LocalTensor<int32_t> &aIdxAll,
215+ const LocalTensor<int32_t> &bsIdxAll, const LocalTensor<int32_t> &kIdxAll)
216+ {
217+ for (int64_t indicesIndex = 0; indicesIndex < curLoopElements; indicesIndex++) {
218+ int64_t aIndices = aIdxAll.GetValue(indicesIndex); // 三个下标均取自上面算好的向量分道
219+ int64_t bsIndices = bsIdxAll.GetValue(indicesIndex);
220+ int64_t kIndices = kIdxAll.GetValue(indicesIndex);
221+ int32_t sessionIndices = 0;
222+ int32_t microbatchIndices = 0;
223+ if constexpr (isScanFlag == false) {
224+ sessionIndices = sessionIdsLocal_.GetValue(aIndices);
225+ microbatchIndices = microBatchIdsLocal_.GetValue(aIndices);
226+ } else {
227+ sessionIndices = aIndices;
228+ microbatchIndices = curMicroBatchID;
229+ }
230+ 
231+ outAllLocal.SetValue(indicesIndex, sessionIndices);
232+ outAllLocal.SetValue(PER_LOOP_ROWS * VAR_MICRO_BATCH_IDX + indicesIndex, microbatchIndices);
233+ outAllLocal.SetValue(PER_LOOP_ROWS * VAR_TOKEN_IDX + indicesIndex, bsIndices);
234+ outAllLocal.SetValue(PER_LOOP_ROWS * VAR_EXPERT_OFFSETS_IDX + indicesIndex, kIndices);
235+ 
236+ for (int64_t hBlock = 0; hBlock < hBlocks_; hBlock++) {
237+ int64_t hStart = hBlock * maxBlockSize_;
238+ int64_t hSize = (hBlock == hBlocks_ - 1) ? lastHBlockSize_ : maxBlockSize_;
239+ int64_t globalXOffset = sessionIndices * strideSession + microbatchIndices * strideMicroBatch +
240+ bsIndices * strideBs + kIndices * strideK + hStart;
241+ 
242+ bool isLastBlock = (hBlock == hBlocks_ - 1);
243+ CopyXIn(globalXOffset, hSize, indicesIndex, outAllLocal[PER_LOOP_ROWS * VAR_NUM], isLastBlock);
244+ int64_t outputOffset =
245+ (indicesIndex + currentOuterStart) * contextInfo_->H * tokenDtypeSize_ + hStart * tokenDtypeSize_;
246+ CopyXOut(outputOffset, hSize);
247+ }
248+ }
249+ }
250+ 
251+private:
252+ // 把扁平 gatherIdx 拆成 (aIdx, bsIdx, kIdx) 三路分道:整块向量运算,
253+ // 避免逐 token 做两次整数除法。三路结果供下方标量循环逐元素取用。
254+ __aicore__ inline void DecomposeIndices(const LocalTensor<int32_t> &expertIdxLocal, int64_t bskProduct,
255+ int64_t curLoopElements)
256+ {
257+ // ---- 索引换算(向量):把逐 token 的两次整数除法整块算完 ----
258+ // aIdx = gatherIdx / bskProduct;bsIdx = 余数 / K;kIdx = 余数 % K。
259+ // gatherIdx < Y <= 2^22,fp32 精确表示 2^24 内整数,商与余数无误差;商恒非负。
260+ LocalTensor<float> lanes = idxCalcBuf_.Get<float>();
261+ LocalTensor<float> srcF = lanes[PER_LOOP_ROWS * IDX_LANE_SRC];
262+ LocalTensor<float> qaF = lanes[PER_LOOP_ROWS * IDX_LANE_QA];
263+ LocalTensor<float> aF = lanes[PER_LOOP_ROWS * IDX_LANE_AF];
264+ LocalTensor<float> remF = lanes[PER_LOOP_ROWS * IDX_LANE_REM];
265+ LocalTensor<float> qbF = lanes[PER_LOOP_ROWS * IDX_LANE_QB];
266+ LocalTensor<float> bF = lanes[PER_LOOP_ROWS * IDX_LANE_BF];
267+ LocalTensor<float> kF = lanes[PER_LOOP_ROWS * IDX_LANE_KF];
268+ LocalTensor<int32_t> ints = idxIntBuf_.Get<int32_t>();
269+ LocalTensor<int32_t> aIdxAll = ints[PER_LOOP_ROWS * IDX_LANE_AI];
270+ LocalTensor<int32_t> bsIdxAll = ints[PER_LOOP_ROWS * IDX_LANE_BI];
271+ LocalTensor<int32_t> kIdxAll = ints[PER_LOOP_ROWS * IDX_LANE_KI];
272+ // 除法用倒数乘法(硬件无整数向量除法),但 1/n 与乘积各有一次舍入,恰好整除处
273+ // floor 会掉一档(实测 BS*K=189 时 a 偏 -1、bs 偏 +BS,k 因两处偏移抵消反而不变)。
274+ // 故每次除法后都按余数做一次 ±1 修正:余数是整数且落在 (-n, 2n),
275+ // 0/1 指示可以纯算术拿到,不需要比较掩码与 Select:
276+ // hi = min(max(rem - n + 1, 0), 1) —— rem >= n 时为 1
277+ // lo = min(max(-rem, 0), 1) —— rem < 0 时为 1
278+ // 修正量 corr = hi - lo ∈ {-1, 0, 1}。全程整数值,fp32 精确(gatherIdx < 2^24)。
279+ const float bskF = static_cast<float>(bskProduct);
280+ const float kNumF = static_cast<float>(contextInfo_->K);
281+ Cast(srcF, expertIdxLocal, RoundMode::CAST_ROUND, curLoopElements);
282+ PipeBarrier<PIPE_V>();
283+ Muls(qaF, srcF, static_cast<float>(1.0f / bskF), curLoopElements);
284+ PipeBarrier<PIPE_V>();
285+ Cast(aIdxAll, qaF, RoundMode::CAST_FLOOR, curLoopElements);
286+ PipeBarrier<PIPE_V>();
287+ Cast(aF, aIdxAll, RoundMode::CAST_ROUND, curLoopElements);
288+ PipeBarrier<PIPE_V>();
289+ Muls(remF, aF, -bskF, curLoopElements);
290+ PipeBarrier<PIPE_V>();
291+ Add(remF, remF, srcF, curLoopElements); // rem = gatherIdx - aIdx*bskProduct
292+ PipeBarrier<PIPE_V>();
293+ CorrectQuotient(aF, remF, qaF, qbF, bskF, curLoopElements);
294+ Cast(aIdxAll, aF, RoundMode::CAST_RINT, curLoopElements);
295+ PipeBarrier<PIPE_V>();
296+ 
297+ Muls(qbF, remF, static_cast<float>(1.0f / kNumF), curLoopElements);
298+ PipeBarrier<PIPE_V>();
299+ Cast(bsIdxAll, qbF, RoundMode::CAST_FLOOR, curLoopElements);
300+ PipeBarrier<PIPE_V>();
301+ Cast(bF, bsIdxAll, RoundMode::CAST_ROUND, curLoopElements);
302+ PipeBarrier<PIPE_V>();
303+ Muls(kF, bF, -kNumF, curLoopElements);
304+ PipeBarrier<PIPE_V>();
305+ Add(kF, remF, kF, curLoopElements); // kIdx = rem - bsIdx*K
306+ PipeBarrier<PIPE_V>();
307+ CorrectQuotient(bF, kF, qaF, qbF, kNumF, curLoopElements);
308+ Cast(bsIdxAll, bF, RoundMode::CAST_RINT, curLoopElements);
309+ PipeBarrier<PIPE_V>();
310+ Cast(kIdxAll, kF, RoundMode::CAST_RINT, curLoopElements);
311+ SetWaitFlag<HardEvent::V_S>(HardEvent::V_S);
312+ }
313+ 
314+ // 倒数乘法得到的商可能差 1,用余数把商与余数一起拉回正确值。
315+ // quot/rem 传入传出均为整数值的 fp32;t0/t1 是临时分道,调用后内容作废。
316+ __aicore__ inline void CorrectQuotient(const LocalTensor<float> &quot, const LocalTensor<float> &rem,
317+ const LocalTensor<float> &t0, const LocalTensor<float> &t1, float divisor,
318+ int64_t count)
319+ {
320+ Adds(t0, rem, 1.0f - divisor, count); // rem - divisor + 1
321+ PipeBarrier<PIPE_V>();
322+ Maxs(t0, t0, 0.0f, count);
323+ PipeBarrier<PIPE_V>();
324+ Mins(t0, t0, 1.0f, count); // hi = (rem >= divisor)
325+ PipeBarrier<PIPE_V>();
326+ Muls(t1, rem, -1.0f, count);
327+ PipeBarrier<PIPE_V>();
328+ Maxs(t1, t1, 0.0f, count);
329+ PipeBarrier<PIPE_V>();
330+ Mins(t1, t1, 1.0f, count); // lo = (rem < 0)
331+ PipeBarrier<PIPE_V>();
332+ Sub(t0, t0, t1, count); // corr = hi - lo
333+ PipeBarrier<PIPE_V>();
334+ Add(quot, t0, quot, count);
335+ PipeBarrier<PIPE_V>();
336+ Muls(t1, t0, -divisor, count);
337+ PipeBarrier<PIPE_V>();
338+ Add(rem, t1, rem, count);
339+ PipeBarrier<PIPE_V>();
340+ }
341+ 
342+ __aicore__ inline void CopyAllLocalOut(int64_t allLocalOffset, int64_t copyLength)
343+ {
344+ LocalTensor<int32_t> outAllLocal = outQueALL_.DeQue<int32_t>();
345+ 
346+ DataCopyExtParams copyParams2{1, static_cast<uint32_t>(copyLength * sizeof(int32_t)), 0, 0, 0};
347+ DataCopyPad(sessionIdsOutGm_[allLocalOffset], outAllLocal, copyParams2);
348+ DataCopyPad(microBatchIdsOutGm_[allLocalOffset], outAllLocal[PER_LOOP_ROWS * VAR_MICRO_BATCH_IDX], copyParams2);
349+ DataCopyPad(tokenIdsOutGm_[allLocalOffset], outAllLocal[PER_LOOP_ROWS * VAR_TOKEN_IDX], copyParams2);
350+ DataCopyPad(expertOffsetsOutGm_[allLocalOffset], outAllLocal[PER_LOOP_ROWS * VAR_EXPERT_OFFSETS_IDX],
351+ copyParams2);
352+ 
353+ if (contextInfo_->tokenDtype == TOKEN_KIND_TWO) {
354+ LocalTensor<int32_t> srcOffsetLocal =
355+ outAllLocal[PER_LOOP_ROWS * VAR_DYNAMIC_SCALE].template ReinterpretCast<int32_t>();
356+ LocalTensor<float> dynamicScaleLocalFp32 =
357+ outAllLocal[PER_LOOP_ROWS * VAR_NUM].template ReinterpretCast<float>();
358+ ArithProgression<int32_t>(srcOffsetLocal, 0, BLOCK_SIZE, copyLength);
359+ PipeBarrier<PIPE_V>();
360+ Gather(dynamicScaleLocalFp32, dynamicScaleLocalFp32, srcOffsetLocal.template ReinterpretCast<uint32_t>(), 0,
361+ copyLength);
362+ SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
363+ DataCopyPad(dynamicScaleOutGm_[allLocalOffset], dynamicScaleLocalFp32, copyParams2);
364+ }
365+ outQueALL_.FreeTensor(outAllLocal);
366+ }
367+ 
368+ __aicore__ inline void CopyXIn(int64_t xSrcOffset, int64_t curLoopCols, int64_t indicesIndex,
369+ const LocalTensor<int32_t> &dynamicScaleLocal, bool isLastBlock)
370+ {
371+ LocalTensor<int8_t> xLocal = inQueueX_.AllocTensor<int8_t>();
372+ uint32_t copySize = curLoopCols * tokenDtypeSize_;
373+ DataCopyExtParams copyParams0{1, copySize, 0, 0, 0};
374+ DataCopyPadExtParams<int8_t> padParams0{false, 0, 0, 0};
375+ DataCopyPad(xLocal, tokenDataBufGm_[xSrcOffset], copyParams0, padParams0);
376+ 
377+ if (isLastBlock && contextInfo_->tokenDtype == TOKEN_KIND_TWO) {
378+ DataCopyExtParams copyParams1{1, static_cast<uint32_t>(sizeof(float)), 0, 0, 0};
379+ // 8 blocks
380+ LocalTensor<int8_t> dynScaleT = dynamicScaleLocal[indicesIndex * 8].template ReinterpretCast<int8_t>();
381+ DataCopyPad(dynScaleT, tokenDataBufGm_[xSrcOffset + copySize], copyParams1, padParams0);
382+ }
383+ 
384+ inQueueX_.EnQue(xLocal);
385+ }
386+ 
387+ __aicore__ inline void CopyXOut(int64_t xDstOffset, int64_t curLoopCols)
388+ {
389+ LocalTensor<int8_t> xLocal = inQueueX_.DeQue<int8_t>();
390+ 
391+ DataCopyExtParams copyParams2{1, static_cast<uint32_t>(curLoopCols * tokenDtypeSize_), 0, 0, 0};
392+ DataCopyPad(yOutGm_[xDstOffset], xLocal, copyParams2);
393+ 
394+ inQueueX_.FreeTensor(xLocal);
395+ }
396+ 
397+private:
398+ static constexpr uint32_t BLOCK_SIZE = 32;
399+ static constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue
400+ static constexpr int32_t PER_LOOP_ROWS = 128; // tensor num for each queue
401+ static constexpr int64_t TOKEN_KIND_TWO = 2;
402+ 
403+ static constexpr int32_t VAR_NUM = 5; // session_ids, micro_batch_ids, token_ids, expert_offsets
404+ static constexpr int32_t VAR_SESSION_IDX = 0;
405+ static constexpr int32_t VAR_MICRO_BATCH_IDX = 1;
406+ static constexpr int32_t VAR_TOKEN_IDX = 2;
407+ static constexpr int32_t VAR_EXPERT_OFFSETS_IDX = 3;
408+ static constexpr int32_t VAR_DYNAMIC_SCALE = 4;
409+ 
410+ TQueBind<TPosition::VECIN, TPosition::VECOUT, BUFFER_NUM> inQueueX_;
411+ TQue<TPosition::VECOUT, BUFFER_NUM> outQueALL_;
412+ TQue<TPosition::VECIN, BUFFER_NUM> expertIdxQue_;
413+ TBuf<TPosition::VECIN> tmpBuffer_;
414+ TBuf<TPosition::VECCALC> idxCalcBuf_; // fp32 中转分道
415+ TBuf<TPosition::VECCALC> idxIntBuf_; // int32 结果分道(类型独立,避免 ReinterpretCast)
416+ GlobalTensor<int8_t> tokenDataBufGm_; // 这里的token_data_buf 存储的是[A,M,BS,K,HS]
417+ GlobalTensor<int32_t> expertIdxGm_;
418+ GlobalTensor<int32_t> sessionIdsInGm_;
419+ GlobalTensor<int32_t> microBatchIdsInGm_;
420+ 
421+ GlobalTensor<int8_t> yOutGm_;
422+ GlobalTensor<int32_t> sessionIdsOutGm_;
423+ GlobalTensor<int32_t> microBatchIdsOutGm_;
424+ GlobalTensor<int32_t> tokenIdsOutGm_;
425+ GlobalTensor<int32_t> expertOffsetsOutGm_;
426+ GlobalTensor<float> dynamicScaleOutGm_;
427+ 
428+ LocalTensor<int32_t> sessionIdsLocal_;
429+ LocalTensor<int32_t> microBatchIdsLocal_;
430+ 
431+ const ScheduleContextInfo *contextInfo_ = nullptr;
432+ 
433+ int64_t needCoreNum_ = 0;
434+ int64_t lastLoopRows_ = 0;
435+ int64_t rowLoops_ = 0;
436+ int64_t tokenDtypeSize_ = 0;
437+ int64_t sessionNumBlockAlign_ = 0;
438+ 
439+ int64_t maxBlockSize_ = 0;
440+ int64_t hBlocks_ = 0;
441+ int64_t lastHBlockSize_ = 0;
442+ uint32_t curMicroBatchID = 0;
443+ int64_t BsKPaddingCount = 0;
444+};
445+ 
446+} // namespace FfnWbBatchingArch35
447+#endif // OP_KERNEL_ARCH35_FFN_WB_A5_GATHER_H
Affn/ffn_worker_batching/op_kernel/arch35/ffn_wb_a5_group_list.h+377-0
@@ -0,0 +1,377 @@
1+/**
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
3+ *
4+ * Licensed under the Apache License, Version 2.0 (the "License");
5+ * you may not use this file except in compliance with the License.
6+ * You may obtain a copy of the License at
7+ *
8+ * http://www.apache.org/licenses/LICENSE-2.0
9+ *
10+ * Unless required by applicable law or agreed to in writing, software
11+ * distributed under the License is distributed on an "AS IS" BASIS,
12+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+ * See the License for the specific language governing permissions and
14+ * limitations under the License.
15+ */
16+ 
17+/*!
18+ * \file ffn_wb_a5_group_list.h
19+ * \brief arch35(A5) group_list:从有序 expert_id 直接压出 [expert_id, tokenNum] 稠密行。
20+ *
21+ * 序列已按 expert_id 升序,故"与前一个不同"即某专家的首元素、"与后一个不同"即其末元素。
22+ * 两张掩码各做一次 GatherMask 压缩,得到**等长**的首下标列与末下标列,逐元素相减即每专家
23+ * token 数。全程向量:不按 expert_id 散射写,也不需要 GM 直方图与跨核原子累加。
24+ *
25+ * 为何不用 SIMT 按 expert_id 散射计数:SIMT 访问 __local_mem__ 被限制在 UB 的低 8KB
26+ * (固定值,与 asc_vf_call 的 dim3 线程数无关,实测把线程数减半阈值不动),按 expert_id
27+ * 写 slot[cur] 在专家数超过 2048 时越界;越界写静默生效,会踩坏同一 UB 上的其它数据
28+ * (曾表现为编译器溢出到栈上的 totalValid 被清零,后续 gather 整段被跳过而硬件不报异常)。
29+ *
30+ * 分块由运行时 UB 反推(见 Init 中的 chunkElements_),不写死块长;块间用一个部分游程
31+ * (pendingId/pendingCnt)承接,故任意长度的游程都能跨块正确合并。
32+ *
33+ * 写出严格走整段搬运:行在 UB 拼好后按 host 给的块长 DataCopyPad 一次搬出,未使用的专家行
34+ * 用 Duplicate 填 0 随块带出。**不对 GM 做标量 SetValue**——历史缺陷正是"逐行 SetValue +
35+ * 末尾一次 SINGLE_CACHE_LINE 刷回"导致只有首条 cache line(4 行)落盘。
36+ */
37+#ifndef OP_KERNEL_ARCH35_FFN_WB_A5_GROUP_LIST_H
38+#define OP_KERNEL_ARCH35_FFN_WB_A5_GROUP_LIST_H
39+#include "ffn_wb_a5_context.h"
40+ 
41+namespace FfnWbBatchingArch35 {
42+using namespace AscendC;
43+ 
44+constexpr int64_t ONE_REPEAT_COMPARE_NUM = 64; // CompareScalar 的 repeat 粒度(ISA)
45+constexpr int64_t ONE_REPEAT_BLOCKS = 8; // 一个 repeat 覆盖 8 个 32B 块(ISA)
46+constexpr int64_t ONE_BYTE_BITS = 8;
47+constexpr int32_t SENTINEL_ID = -1; // 哨兵:有效 expert_id 恒 >= 0,故它与谁都不等
48+ 
49+// 单元素在 UB 上的占用:cur/prev/next/idx/diff/diffF/outStart/outEnd/outId/cnt 十路 int32,
50+// 外加首尾两张比较掩码(各 1 bit/元素)。分块大小由 UB 除以它得到,不是拍出来的常数。
51+constexpr int64_t CHUNK_VEC_NUM = 10;
52+constexpr int64_t CHUNK_MASK_NUM = 2;
53+constexpr int64_t CHUNK_BUF_NUM = CHUNK_VEC_NUM + CHUNK_MASK_NUM + 1; // +1:行拼装区
54+ 
55+struct A5GroupListParam {
56+ int64_t rowsPerLoop = 0; // 写出时每块拼多少行(每行 16B)
57+};
58+ 
59+// 游程编码用到的 UB 分道集合:一次取好,三个子步骤共用,免得每个函数各取一遍。
60+struct A5RunLanes {
61+ LocalTensor<int32_t> cur; // ids[begin, begin+len)
62+ LocalTensor<int32_t> prev; // ids[begin-1, begin+len-1)
63+ LocalTensor<int32_t> next; // ids[begin+1, begin+len]
64+ LocalTensor<int32_t> idx; // 块内局部下标 0,1,2,...
65+ LocalTensor<int32_t> diff; // 相邻差(整型)
66+ LocalTensor<float> diffF; // 相邻差(浮点,供比较指令)
67+ LocalTensor<int32_t> outStart; // 压缩后的游程首下标
68+ LocalTensor<int32_t> outEnd; // 压缩后的游程末下标
69+ LocalTensor<int32_t> outId; // 压缩后的游程 expert_id
70+ LocalTensor<int32_t> cnt; // 每游程 token 数
71+ LocalTensor<uint32_t> maskS; // 游程起点掩码
72+ LocalTensor<uint32_t> maskE; // 游程终点掩码
73+};
74+ 
75+class FfnWbA5GroupList {
76+public:
77+ __aicore__ inline FfnWbA5GroupList(){};
78+ 
79+ // sortedIdsWs:有序 expert_id(长度为排序后的有效数);groupList:输出 [expertNum, 2]
80+ __aicore__ inline void Init(GM_ADDR sortedIdsWs, GM_ADDR groupList, const A5GroupListParam &param,
81+ const ScheduleContextInfo *ctx, TPipe *pipe)
82+ {
83+ param_ = param;
84+ pipe_ = pipe;
85+ const int64_t subNum = GetTaskRation() > 0 ? GetTaskRation() : 1;
86+ vecId_ = GetBlockIdx() * subNum + GetSubBlockIdx();
87+ expertNum_ = static_cast<int64_t>(ctx->expertNum);
88+ totalLen_ = ctx->validGatherIdxLength;
89+ // 游程压缩是串行语义(相邻元素比较跨越切片边界),由 0 号核一趟流式做完;
90+ // 其余核不参与,也就不占 UB——它们直接进入下一相位。
91+ if (vecId_ != 0) {
92+ return;
93+ }
94+ 
95+ idsGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(sortedIdsWs), (totalLen_ > 0) ? totalLen_ : 1);
96+ groupListGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(groupList), expertNum_ * NUM_TWO);
97+ 
98+ // 分块长度由运行时 UB 反推:先扣掉行拼装区与各缓冲的块对齐余量,再按每元素占用摊分,
99+ // 最后向下取整到比较指令的 repeat 粒度。
100+ const int64_t rowBytes = param_.rowsPerLoop * NUM_TWO * static_cast<int64_t>(sizeof(int64_t));
101+ const int64_t reserved = rowBytes + CHUNK_BUF_NUM * ONE_BLK_SIZE;
102+ const int64_t avail = (ctx->ubSize > reserved) ? (ctx->ubSize - reserved) : 0;
103+ const int64_t bytesPerElem = CHUNK_VEC_NUM * static_cast<int64_t>(sizeof(int32_t)) + CHUNK_MASK_NUM;
104+ int64_t chunk = avail / bytesPerElem / ONE_REPEAT_COMPARE_NUM * ONE_REPEAT_COMPARE_NUM;
105+ if (chunk < ONE_REPEAT_COMPARE_NUM) {
106+ chunk = ONE_REPEAT_COMPARE_NUM;
107+ }
108+ // 数据装得下就不必开满,按实际长度收窄(仍保持 repeat 粒度)。
109+ const int64_t needed = Ceil((totalLen_ > 0) ? totalLen_ : 1, ONE_REPEAT_COMPARE_NUM) * ONE_REPEAT_COMPARE_NUM;
110+ chunkElements_ = (chunk > needed) ? needed : chunk;
111+ 
112+ const int64_t vecBytes = chunkElements_ * static_cast<int64_t>(sizeof(int32_t)) + ONE_BLK_SIZE;
113+ pipe_->InitBuffer(curBuf_, vecBytes);
114+ pipe_->InitBuffer(prevBuf_, vecBytes);
115+ pipe_->InitBuffer(nextBuf_, vecBytes);
116+ pipe_->InitBuffer(idxBuf_, vecBytes);
117+ pipe_->InitBuffer(diffBuf_, vecBytes);
118+ pipe_->InitBuffer(diffFBuf_, vecBytes);
119+ pipe_->InitBuffer(outStartBuf_, vecBytes);
120+ pipe_->InitBuffer(outEndBuf_, vecBytes);
121+ pipe_->InitBuffer(outIdBuf_, vecBytes);
122+ pipe_->InitBuffer(cntBuf_, vecBytes);
123+ const int64_t maskBytes = chunkElements_ / ONE_BYTE_BITS + ONE_BLK_SIZE;
124+ pipe_->InitBuffer(maskSBuf_, maskBytes);
125+ pipe_->InitBuffer(maskEBuf_, maskBytes);
126+ pipe_->InitBuffer(rowBuf_, rowBytes + ONE_BLK_SIZE);
127+ 
128+ // 块内局部下标序列与块无关,只生成一次;计数用差值,局部下标与全局下标同解。
129+ ArithProgression<int32_t>(idxBuf_.Get<int32_t>(), 0, 1, static_cast<int32_t>(chunkElements_));
130+ PipeBarrier<PIPE_V>();
131+ }
132+ 
133+ __aicore__ inline void Process()
134+ {
135+ if (vecId_ != 0) {
136+ return;
137+ }
138+ LocalTensor<int64_t> rows = rowBuf_.Get<int64_t>();
139+ LocalTensor<int32_t> rows32 = rows.template ReinterpretCast<int32_t>();
140+ ClearRows(rows32);
141+ 
142+ int64_t outIdx = 0; // 已搬出的行数
143+ int64_t fill = 0; // 当前块已拼的行数
144+ if (totalLen_ > 0) {
145+ EncodeRuns(rows, outIdx, fill);
146+ }
147+ FlushRows(rows, outIdx, fill);
148+ // 未用到的专家:整块 [0,0] 补齐到 expertNum 行(块在上一次搬出后已清零)。
149+ while (outIdx < expertNum_) {
150+ const int64_t n = Min(param_.rowsPerLoop, expertNum_ - outIdx);
151+ SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
152+ DataCopyExtParams cpOut{static_cast<uint16_t>(1), static_cast<uint32_t>(n * NUM_TWO * sizeof(int64_t)), 0,
153+ 0, 0};
154+ DataCopyPad(groupListGm_[outIdx * NUM_TWO], rows, cpOut);
155+ SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
156+ outIdx += n;
157+ }
158+ }
159+ 
160+private:
161+ // 流式游程编码:逐块比较"与前一个/与后一个是否不同",压出完整游程的首尾下标。
162+ // 一次取好全部分道视图。
163+ __aicore__ inline A5RunLanes MakeLanes()
164+ {
165+ A5RunLanes ln;
166+ ln.cur = curBuf_.Get<int32_t>();
167+ ln.prev = prevBuf_.Get<int32_t>();
168+ ln.next = nextBuf_.Get<int32_t>();
169+ ln.idx = idxBuf_.Get<int32_t>();
170+ ln.diff = diffBuf_.Get<int32_t>();
171+ ln.diffF = diffFBuf_.Get<float>();
172+ ln.outStart = outStartBuf_.Get<int32_t>();
173+ ln.outEnd = outEndBuf_.Get<int32_t>();
174+ ln.outId = outIdBuf_.Get<int32_t>();
175+ ln.cnt = cntBuf_.Get<int32_t>();
176+ ln.maskS = maskSBuf_.Get<uint32_t>();
177+ ln.maskE = maskEBuf_.Get<uint32_t>();
178+ return ln;
179+ }
180+ 
181+ // 载入一块的 cur/prev/next 三路,并按边界置哨兵:块首一律记为游程起点(缺的那段由
182+ // pendingCnt 承接),全局末元素一律记为游程终点。
183+ __aicore__ inline void LoadChunk(A5RunLanes &ln, const DataCopyPadExtParams<int32_t> &pad, int64_t begin,
184+ int64_t len)
185+ {
186+ DataCopyExtParams cpLen{static_cast<uint16_t>(1), static_cast<uint32_t>(len * sizeof(int32_t)), 0, 0, 0};
187+ DataCopyPad(ln.cur, idsGm_[begin], cpLen, pad);
188+ DataCopyPad(ln.prev, idsGm_[begin - 1], cpLen, pad);
189+ const int64_t nextLen = Min(len, totalLen_ - begin - 1);
190+ if (nextLen > 0) {
191+ DataCopyExtParams cpNext{static_cast<uint16_t>(1), static_cast<uint32_t>(nextLen * sizeof(int32_t)), 0, 0,
192+ 0};
193+ DataCopyPad(ln.next, idsGm_[begin + 1], cpNext, pad);
194+ }
195+ SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
196+ // 块首一律记为游程起点:它若处在游程中间,缺的那一段正由 pendingCnt 承接。
197+ ln.prev.SetValue(0, SENTINEL_ID);
198+ if (nextLen < len) {
199+ ln.next.SetValue(len - 1, SENTINEL_ID); // 全局末元素必是游程终点
200+ }
201+ SetWaitFlag<HardEvent::S_V>(HardEvent::S_V);
202+ }
203+ 
204+ // 出首尾两张掩码并压缩成等长的三列,返回本块内完整游程数。
205+ __aicore__ inline int64_t BuildAndCompact(A5RunLanes &ln, int64_t len, uint64_t &startNum, uint64_t &endNum)
206+ {
207+ const int64_t cmpLen = Ceil(len, ONE_REPEAT_COMPARE_NUM) * ONE_REPEAT_COMPARE_NUM;
208+ // 比较指令只吃浮点,故差值转 float 再与 0 比;差值幅度不超过专家数上界,转换精确。
209+ Sub(ln.diff, ln.cur, ln.prev, static_cast<int32_t>(len));
210+ PipeBarrier<PIPE_V>();
211+ Cast(ln.diffF, ln.diff, RoundMode::CAST_ROUND, static_cast<int32_t>(len));
212+ PipeBarrier<PIPE_V>();
213+ CompareScalar(ln.maskS.template ReinterpretCast<uint8_t>(), ln.diffF, static_cast<float>(0), CMPMODE::NE,
214+ static_cast<int32_t>(cmpLen));
215+ PipeBarrier<PIPE_V>();
216+ Sub(ln.diff, ln.next, ln.cur, static_cast<int32_t>(len));
217+ PipeBarrier<PIPE_V>();
218+ Cast(ln.diffF, ln.diff, RoundMode::CAST_ROUND, static_cast<int32_t>(len));
219+ PipeBarrier<PIPE_V>();
220+ CompareScalar(ln.maskE.template ReinterpretCast<uint8_t>(), ln.diffF, static_cast<float>(0), CMPMODE::NE,
221+ static_cast<int32_t>(cmpLen));
222+ PipeBarrier<PIPE_V>();
223+ 
224+ GatherMaskParams gmp;
225+ gmp.repeatTimes = 1;
226+ gmp.src0BlockStride = 1;
227+ gmp.src0RepeatStride = ONE_REPEAT_BLOCKS;
228+ gmp.src1RepeatStride = 0;
229+ uint64_t idNum = 0;
230+ GatherMask(ln.outStart, ln.idx, ln.maskS, true, static_cast<uint32_t>(len), gmp, startNum);
231+ GatherMask(ln.outEnd, ln.idx, ln.maskE, true, static_cast<uint32_t>(len), gmp, endNum);
232+ GatherMask(ln.outId, ln.cur, ln.maskE, true, static_cast<uint32_t>(len), gmp, idNum);
233+ PipeBarrier<PIPE_V>();
234+ 
235+ // 在本块结束的游程即完整游程,其首尾下标在两列中一一对应(块首已强制成起点)。
236+ const int64_t runs = static_cast<int64_t>(endNum);
237+ if (runs > 0) {
238+ Sub(ln.cnt, ln.outEnd, ln.outStart, static_cast<int32_t>(runs));
239+ PipeBarrier<PIPE_V>();
240+ Adds(ln.cnt, ln.cnt, 1, static_cast<int32_t>(runs));
241+ PipeBarrier<PIPE_V>();
242+ }
243+ SetWaitFlag<HardEvent::V_S>(HardEvent::V_S);
244+ return static_cast<int64_t>(endNum);
245+ }
246+ 
247+ __aicore__ inline void EncodeRuns(const LocalTensor<int64_t> &rows, int64_t &outIdx, int64_t &fill)
248+ {
249+ A5RunLanes ln = MakeLanes();
250+ LocalTensor<int32_t> &cur = ln.cur;
251+ LocalTensor<int32_t> &outStart = ln.outStart;
252+ LocalTensor<int32_t> &outId = ln.outId;
253+ LocalTensor<int32_t> &cnt = ln.cnt;
254+ DataCopyPadExtParams<int32_t> pad{false, 0, 0, 0};
255+ 
256+ // 首元素单独起头:此后每块的起点都有前驱,ln.prev 可整块从 ids[begin-1] 读入(免去
257+ // 4 字节偏移的非对齐视图)。它本身必是某游程的首元素,故作为初始的部分游程。
258+ DataCopyExtParams cpOne{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(int32_t)), 0, 0, 0};
259+ DataCopyPad(cur, idsGm_, cpOne, pad);
260+ SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
261+ int64_t pendingId = static_cast<int64_t>(cur.GetValue(0));
262+ int64_t pendingCnt = 1;
263+ 
264+ int64_t begin = 1;
265+ while (begin < totalLen_) {
266+ const int64_t len = Min(chunkElements_, totalLen_ - begin);
267+ LoadChunk(ln, pad, begin, len);
268+ 
269+ uint64_t startNum = 0;
270+ uint64_t endNum = 0;
271+ const int64_t runs = BuildAndCompact(ln, len, startNum, endNum);
272+ 
273+ // 承接上一块的部分游程:块首若换了专家,说明那段已经完整,先单独成行。
274+ // (块中途留下的 pending 必与本块首元素同专家,故此判定只会在起头那一个元素上生效。)
275+ if (pendingCnt > 0 && static_cast<int64_t>(cur.GetValue(0)) != pendingId) {
276+ AppendRow(rows, outIdx, fill, pendingId, pendingCnt);
277+ pendingCnt = 0;
278+ }
279+ if (pendingCnt == 0) {
280+ pendingId = static_cast<int64_t>(cur.GetValue(0)); // 本块首元素开启的游程
281+ }
282+ 
283+ if (runs == 0) {
284+ pendingCnt += len; // 整块仍落在同一个游程内,继续往后并
285+ } else {
286+ for (int64_t j = 0; j < runs; j++) {
287+ int64_t num = static_cast<int64_t>(cnt.GetValue(j));
288+ if (j == 0) {
289+ num += pendingCnt; // 首个游程接上前面块里的部分
290+ }
291+ AppendRow(rows, outIdx, fill, static_cast<int64_t>(outId.GetValue(j)), num);
292+ }
293+ // 块尾若停在游程中间,把这段残长转为下一块的部分游程。
294+ if (static_cast<int64_t>(startNum) > runs) {
295+ pendingCnt = len - static_cast<int64_t>(outStart.GetValue(runs));
296+ pendingId = static_cast<int64_t>(cur.GetValue(len - 1));
297+ } else {
298+ pendingCnt = 0;
299+ }
300+ }
301+ begin += len;
302+ }
303+ // 全局末元素被哨兵判为游程终点,故循环出来时通常已无残留;
304+ // 唯独 totalLen_ == 1(循环体没进过)时,起头的那个游程要在这里写出。
305+ if (pendingCnt > 0) {
306+ AppendRow(rows, outIdx, fill, pendingId, pendingCnt);
307+ }
308+ }
309+ 
310+ __aicore__ inline void AppendRow(const LocalTensor<int64_t> &rows, int64_t &outIdx, int64_t &fill, int64_t expertId,
311+ int64_t tokenNum)
312+ {
313+ if (outIdx + fill >= expertNum_) {
314+ return; // 游程数不会超过专家数;越界即异常输入,按契约不再写
315+ }
316+ rows.SetValue(fill * NUM_TWO, expertId);
317+ rows.SetValue(fill * NUM_TWO + 1, tokenNum);
318+ fill++;
319+ if (fill == param_.rowsPerLoop) {
320+ FlushRows(rows, outIdx, fill);
321+ }
322+ }
323+ 
324+ __aicore__ inline void FlushRows(const LocalTensor<int64_t> &rows, int64_t &outIdx, int64_t &fill)
325+ {
326+ if (fill <= 0) {
327+ return;
328+ }
329+ SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
330+ DataCopyExtParams cp{static_cast<uint16_t>(1), static_cast<uint32_t>(fill * NUM_TWO * sizeof(int64_t)), 0, 0,
331+ 0};
332+ DataCopyPad(groupListGm_[outIdx * NUM_TWO], rows, cp);
333+ SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
334+ outIdx += fill;
335+ fill = 0;
336+ // 复位缓冲:不复位的话,下一块未填满的行会带出上一块的残留值。
337+ ClearRows(rows.template ReinterpretCast<int32_t>());
338+ }
339+ 
340+ // 整块清零:起点为缓冲区首地址,满足向量指令的对齐要求。
341+ __aicore__ inline void ClearRows(const LocalTensor<int32_t> &rows32)
342+ {
343+ // 清零是向量写:既要等前面的标量填充,也要等上一块的搬出真正读完。
344+ SetWaitFlag<HardEvent::MTE3_V>(HardEvent::MTE3_V);
345+ SetWaitFlag<HardEvent::S_V>(HardEvent::S_V);
346+ Duplicate<int32_t>(rows32, 0, static_cast<int32_t>(param_.rowsPerLoop * NUM_FOUR));
347+ SetWaitFlag<HardEvent::V_S>(HardEvent::V_S);
348+ }
349+ 
350+ A5GroupListParam param_;
351+ TPipe *pipe_ = nullptr;
352+ 
353+ GlobalTensor<int32_t> idsGm_;
354+ GlobalTensor<int64_t> groupListGm_;
355+ 
356+ TBuf<TPosition::VECCALC> curBuf_;
357+ TBuf<TPosition::VECCALC> prevBuf_;
358+ TBuf<TPosition::VECCALC> nextBuf_;
359+ TBuf<TPosition::VECCALC> idxBuf_;
360+ TBuf<TPosition::VECCALC> diffBuf_;
361+ TBuf<TPosition::VECCALC> diffFBuf_;
362+ TBuf<TPosition::VECCALC> outStartBuf_;
363+ TBuf<TPosition::VECCALC> outEndBuf_;
364+ TBuf<TPosition::VECCALC> outIdBuf_;
365+ TBuf<TPosition::VECCALC> cntBuf_;
366+ TBuf<TPosition::VECCALC> maskSBuf_;
367+ TBuf<TPosition::VECCALC> maskEBuf_;
368+ TBuf<TPosition::VECCALC> rowBuf_;
369+ 
370+ int64_t vecId_ = 0;
371+ int64_t expertNum_ = 0;
372+ int64_t totalLen_ = 0;
373+ int64_t chunkElements_ = 0;
374+};
375+ 
376+} // namespace FfnWbBatchingArch35
377+#endif // OP_KERNEL_ARCH35_FFN_WB_A5_GROUP_LIST_H
Affn/ffn_worker_batching/op_kernel/arch35/ffn_wb_a5_prepare.h+238-0
@@ -0,0 +1,238 @@
1+/**
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
3+ *
4+ * Licensed under the Apache License, Version 2.0 (the "License");
5+ * you may not use this file except in compliance with the License.
6+ * You may obtain a copy of the License at
7+ *
8+ * http://www.apache.org/licenses/LICENSE-2.0
9+ *
10+ * Unless required by applicable law or agreed to in writing, software
11+ * distributed under the License is distributed on an "AS IS" BASIS,
12+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+ * See the License for the specific language governing permissions and
14+ * limitations under the License.
15+ */
16+ 
17+/*!
18+ * \file ffn_wb_a5_prepare.h
19+ * \brief phase0:等待就绪(RECV) + 把两条路径的 expert_id 归一到同一个扁平缓冲。
20+ */
21+#ifndef OP_KERNEL_ARCH35_FFN_WB_A5_PREPARE_H
22+#define OP_KERNEL_ARCH35_FFN_WB_A5_PREPARE_H
23+#include "ffn_wb_a5_context.h"
24+#include "kernel_operator.h"
25+namespace FfnWbBatchingArch35 {
26+using namespace AscendC;
27+ 
28+// ===================== RECV:等待本 micro batch 就绪 =====================
29+class FfnWbA5RecvWait {
30+public:
31+ __aicore__ inline FfnWbA5RecvWait(){};
32+ 
33+ __aicore__ inline void Init(GM_ADDR schedule_context, GM_ADDR tokenInfoBuf, const ScheduleContextInfo *ctx,
34+ TPipe *pipe)
35+ {
36+ ctx_ = ctx;
37+ pipe_ = pipe;
38+ // FfnDataDesc 每块的 int32 个数:flag + layer_id + expert_ids[BS*K],块数由契约结构给出。
39+ descWords_ = static_cast<int64_t>(sizeof(aicpu::FfnDataDesc)) / static_cast<int64_t>(sizeof(int32_t)) +
40+ static_cast<int64_t>(ctx_->BS) * ctx_->K;
41+ tokenInfoGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(tokenInfoBuf),
42+ static_cast<int64_t>(ctx_->A) * ctx_->M * descWords_);
43+ ctxGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint64_t *>(schedule_context));
44+ 
45+ const int64_t flagWords = static_cast<int64_t>(ctx_->A) * (ONE_BLK_SIZE / sizeof(int32_t));
46+ pipe_->InitBuffer(flagQue_, 1, flagWords * sizeof(int32_t));
47+ pipe_->InitBuffer(workBuf_, flagWords * sizeof(float));
48+ pipe_->InitBuffer(pollBuf_, ONE_BLK_SIZE);
49+ }
50+ 
51+ // 由 0 号核忙等;其余核在调用方的 SyncAll 处等待。
52+ __aicore__ inline void Process()
53+ {
54+ const int64_t subNum = GetTaskRation() > 0 ? GetTaskRation() : 1;
55+ if (GetBlockIdx() * subNum + GetSubBlockIdx() != 0) {
56+ return;
57+ }
58+ const int64_t flagElems = static_cast<int64_t>(ctx_->A) * (ONE_BLK_SIZE / sizeof(int32_t));
59+ const int64_t offset = static_cast<int64_t>(ctx_->curMicroBatchID) * descWords_;
60+ 
61+ while (true) {
62+ LocalTensor<int32_t> flagLocal = flagQue_.AllocTensor<int32_t>();
63+ // 每 session 取一个 flag:块数 A、块长 4B、块间跨度 (M*F-1) 个 int32
64+ DataCopyExtParams cp{
65+ static_cast<uint16_t>(ctx_->A), static_cast<uint32_t>(sizeof(int32_t)),
66+ static_cast<uint32_t>((static_cast<int64_t>(ctx_->M) * descWords_ - 1) * sizeof(int32_t)), 0, 0};
67+ DataCopyPadExtParams<int32_t> pad{
68+ true, 0, static_cast<uint8_t>((ONE_BLK_SIZE - sizeof(int32_t)) / sizeof(int32_t)), 0};
69+ DataCopyPad(flagLocal, tokenInfoGm_[offset], cp, pad);
70+ flagQue_.EnQue(flagLocal);
71+ 
72+ LocalTensor<int32_t> flags = flagQue_.DeQue<int32_t>();
73+ LocalTensor<float> work = workBuf_.Get<float>();
74+ Cast(work, flags, RoundMode::CAST_ROUND, flagElems);
75+ PipeBarrier<PIPE_V>();
76+ ReduceSum<float>(work, work, work, flagElems);
77+ PipeBarrier<PIPE_V>();
78+ SetWaitFlag<HardEvent::V_S>(HardEvent::V_S);
79+ const float readyNum = work.GetValue(0);
80+ flagQue_.FreeTensor(flags);
81+ if (static_cast<uint32_t>(readyNum) >= ctx_->A) {
82+ break;
83+ }
84+ }
85+ 
86+ // 推进轮询下标:偏移取自公共契约结构,回写经 UB 整段搬运。
87+ LocalTensor<uint64_t> pollLocal = pollBuf_.Get<uint64_t>();
88+ pollLocal.SetValue(0, (ctx_->curMicroBatchID + 1) % ctx_->M);
89+ SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
90+ DataCopyExtParams cpPoll{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(uint64_t)), 0, 0, 0};
91+ DataCopyPad(ctxGm_[FFN_WB_CTX_OFFSET(ffn.polling_index) / static_cast<int32_t>(sizeof(uint64_t))], pollLocal,
92+ cpPoll);
93+ SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
94+ }
95+ 
96+private:
97+ const ScheduleContextInfo *ctx_ = nullptr;
98+ TPipe *pipe_ = nullptr;
99+ GlobalTensor<int32_t> tokenInfoGm_;
100+ GlobalTensor<uint64_t> ctxGm_;
101+ TQue<QuePosition::VECIN, 1> flagQue_;
102+ TBuf<TPosition::VECCALC> workBuf_;
103+ TBuf<TPosition::VECCALC> pollBuf_;
104+ int64_t descWords_ = 0;
105+};
106+ 
107+// ===================== expert_id 归一 + 握手回写 =====================
108+// 补位/失效标记:>= sort 的 expertStart_(1000000),排序后落到末尾并被 mask 判据剔除。
109+constexpr int32_t MASK_SENTINEL = 2147483647;
110+ 
111+class FfnWbPrepareArch35 {
112+public:
113+ __aicore__ inline FfnWbPrepareArch35(){};
114+ 
115+ // flatIdsWs:归一后的扁平 expert_id 缓冲(长度 totalLen)。
116+ // rowsPerLoop 由 host 按运行时 UB 容量反推(见 tiling 的 preparePerLoopRows),此处不设任何容量常数。
117+ __aicore__ inline void Init(GM_ADDR flatIdsWs, const ScheduleContextInfo *contextInfo, TPipe *pipe,
118+ int64_t totalLen, int64_t rowsPerLoop)
119+ {
120+ contextInfo_ = contextInfo;
121+ pipe_ = pipe;
122+ totalLen_ = totalLen;
123+ bsk_ = contextInfo_->BS * contextInfo_->K;
124+ F_ = NUM_TWO + bsk_;
125+ // 每 session 行在扁平缓冲中的跨度(RECV 含补位;NORM 无补位时即为 bsk)
126+ rowSpan_ = (contextInfo_->A > 0) ? (totalLen_ / contextInfo_->A) : bsk_;
127+ rowsPerLoop_ = (rowsPerLoop > 0) ? rowsPerLoop : 1;
128+ if (rowsPerLoop_ > contextInfo_->A) {
129+ rowsPerLoop_ = contextInfo_->A;
130+ }
131+ // 按**运行时实际块数**把 session 行分段,各核只处理自己那段:
132+ // 各段写入 flatIds 的区间互不重叠,握手回写也按行分离,故无需跨核同步。
133+ const int64_t coreNum = GetBlockNum();
134+ const int64_t perCore = (contextInfo_->A + coreNum - 1) / coreNum;
135+ rowBegin_ = GetBlockIdx() * perCore;
136+ rowEnd_ = rowBegin_ + perCore;
137+ if (rowEnd_ > contextInfo_->A) {
138+ rowEnd_ = contextInfo_->A;
139+ }
140+ if (rowBegin_ > contextInfo_->A) {
141+ rowBegin_ = contextInfo_->A;
142+ }
143+ if (rowsPerLoop_ > (rowEnd_ - rowBegin_) && (rowEnd_ - rowBegin_) > 0) {
144+ rowsPerLoop_ = rowEnd_ - rowBegin_;
145+ }
146+ flatIdsGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(flatIdsWs), totalLen_);
147+ pipe_->InitBuffer(que_, 1, Align(rowsPerLoop_ * rowSpan_ * sizeof(int32_t), BLOCK_BYTES));
148+ pipe_->InitBuffer(clrBuf_, Align(rowsPerLoop_ * BLOCK_BYTES, BLOCK_BYTES));
149+ }
150+ 
151+ // NORM:expert_ids_buf 已是扁平布局,整段搬到 flatIds(单核即可,量级为 Y 个 int32)。
152+ __aicore__ inline void ProcessNorm(GM_ADDR expertIdsBuf)
153+ {
154+ if (rowEnd_ <= rowBegin_) {
155+ return;
156+ }
157+ GlobalTensor<int32_t> srcGm;
158+ srcGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(expertIdsBuf), totalLen_);
159+ const int64_t perLoopElems = rowsPerLoop_ * rowSpan_;
160+ const int64_t beginElem = rowBegin_ * rowSpan_;
161+ const int64_t endElem = (rowEnd_ * rowSpan_ > totalLen_) ? totalLen_ : rowEnd_ * rowSpan_;
162+ for (int64_t off = beginElem; off < endElem; off += perLoopElems) {
163+ const int64_t n = ((endElem - off) > perLoopElems) ? perLoopElems : (endElem - off);
164+ LocalTensor<int32_t> buf = que_.AllocTensor<int32_t>();
165+ DataCopyExtParams cp{static_cast<uint16_t>(1), static_cast<uint32_t>(n * sizeof(int32_t)), 0, 0, 0};
166+ DataCopyPadExtParams<int32_t> pad{false, 0, 0, 0};
167+ DataCopyPad(buf, srcGm[off], cp, pad);
168+ SetWaitFlag<HardEvent::MTE2_MTE3>(HardEvent::MTE2_MTE3);
169+ DataCopyPad(flatIdsGm_[off], buf, cp);
170+ SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
171+ que_.FreeTensor(buf);
172+ }
173+ }
174+ 
175+ // RECV:逐 session 行跨步取 ids(尾部补 MASK_SENTINEL),并回写 MASK_SENTINEL/0 完成握手。
176+ // 按 rowsPerLoop_ 行一轮处理,轮内 UB 只驻留本轮的 id 区与清零区。
177+ __aicore__ inline void ProcessRecv(GM_ADDR tokenInfoBuf)
178+ {
179+ if (rowEnd_ <= rowBegin_) {
180+ return;
181+ }
182+ GlobalTensor<int32_t> tokenInfoGm;
183+ tokenInfoGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(tokenInfoBuf),
184+ contextInfo_->A * contextInfo_->M * F_);
185+ const int64_t base = contextInfo_->curMicroBatchID * F_;
186+ const int64_t rowStride = contextInfo_->M * F_; // 相邻 session 在 token_info 中的跨度
187+ LocalTensor<int32_t> clr = clrBuf_.Get<int32_t>();
188+ 
189+ for (int64_t r0 = rowBegin_; r0 < rowEnd_; r0 += rowsPerLoop_) {
190+ const int64_t rows = ((rowEnd_ - r0) > rowsPerLoop_) ? rowsPerLoop_ : (rowEnd_ - r0);
191+ LocalTensor<int32_t> buf = que_.AllocTensor<int32_t>();
192+ 
193+ // 取 ids:每块 bsk 个 int32,块间跳过 flag/layer 与其余 micro batch;
194+ // 右侧补位由 DataCopyPad 按 BsKPaddingCount 填 MASK_SENTINEL,使每行在 UB 中占 rowSpan_
195+ DataCopyExtParams inParams{static_cast<uint16_t>(rows), static_cast<uint32_t>(bsk_ * sizeof(int32_t)),
196+ static_cast<uint32_t>((rowStride - bsk_) * sizeof(int32_t)), 0, 0};
197+ DataCopyPadExtParams<int32_t> inPad{true, 0, static_cast<uint8_t>(contextInfo_->BsKPaddingCount),
198+ MASK_SENTINEL};
199+ DataCopyPad(buf, tokenInfoGm[base + r0 * rowStride + NUM_TWO], inParams, inPad);
200+ SetWaitFlag<HardEvent::MTE2_MTE3>(HardEvent::MTE2_MTE3);
201+ 
202+ DataCopyExtParams outParams{static_cast<uint16_t>(1),
203+ static_cast<uint32_t>(rows * rowSpan_ * sizeof(int32_t)), 0, 0, 0};
204+ DataCopyPad(flatIdsGm_[r0 * rowSpan_], buf, outParams);
205+ 
206+ // 握手回写前必须等上面的搬出真正读完 buf:下面要就地把 buf 覆盖成回写内容。
207+ SetWaitFlag<HardEvent::MTE3_V>(HardEvent::MTE3_V);
208+ // 回写值处处相同(ids 全 MASK_SENTINEL、flag 全 0),故按连续块读出即可,与行内跨度无关
209+ Duplicate<int32_t>(buf, MASK_SENTINEL, rows * rowSpan_);
210+ Duplicate<int32_t>(clr, 0, rows * (BLOCK_BYTES / static_cast<int64_t>(sizeof(int32_t))));
211+ SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
212+ DataCopyExtParams clrIds{static_cast<uint16_t>(rows), static_cast<uint32_t>(bsk_ * sizeof(int32_t)), 0,
213+ static_cast<uint32_t>((rowStride - bsk_) * sizeof(int32_t)), 0};
214+ DataCopyExtParams clrFlag{static_cast<uint16_t>(rows), static_cast<uint32_t>(sizeof(int32_t)), 0,
215+ static_cast<uint32_t>((rowStride - 1) * sizeof(int32_t)), 0};
216+ DataCopyPad(tokenInfoGm[base + r0 * rowStride + NUM_TWO], buf, clrIds);
217+ DataCopyPad(tokenInfoGm[base + r0 * rowStride], clr, clrFlag);
218+ SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
219+ que_.FreeTensor(buf);
220+ }
221+ }
222+ 
223+private:
224+ const ScheduleContextInfo *contextInfo_ = nullptr;
225+ TPipe *pipe_ = nullptr;
226+ GlobalTensor<int32_t> flatIdsGm_;
227+ TQue<QuePosition::VECCALC, 1> que_;
228+ TBuf<QuePosition::VECCALC> clrBuf_;
229+ int64_t totalLen_ = 0;
230+ int64_t bsk_ = 0;
231+ int64_t rowSpan_ = 0;
232+ int64_t rowsPerLoop_ = 0;
233+ int64_t rowBegin_ = 0;
234+ int64_t rowEnd_ = 0;
235+ int64_t F_ = 0;
236+};
237+} // namespace FfnWbBatchingArch35
238+#endif // OP_KERNEL_ARCH35_FFN_WB_A5_PREPARE_H
Affn/ffn_worker_batching/op_kernel/arch35/ffn_wb_a5_sort.h+569-0
@@ -0,0 +1,569 @@
1+/**
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
3+ *
4+ * Licensed under the Apache License, Version 2.0 (the "License");
5+ * you may not use this file except in compliance with the License.
6+ * You may obtain a copy of the License at
7+ *
8+ * http://www.apache.org/licenses/LICENSE-2.0
9+ *
10+ * Unless required by applicable law or agreed to in writing, software
11+ * distributed under the License is distributed on an "AS IS" BASIS,
12+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+ * See the License for the specific language governing permissions and
14+ * limitations under the License.
15+ */
16+ 
17+/*!
18+ * \file ffn_wb_a5_sort.h
19+ * \brief phase1:段内排序 → 段间归并 → 拆包,产出有序 expert_id 与 gather_idx。
20+ */
21+#ifndef OP_KERNEL_ARCH35_FFN_WB_A5_SORT_H
22+#define OP_KERNEL_ARCH35_FFN_WB_A5_SORT_H
23+#include "ffn_wb_a5_context.h"
24+namespace FfnWbBatchingArch35 {
25+using namespace AscendC;
26+ 
27+// ===================== 段内排序 =====================
28+// 段内排序的切分参数,全部由 host tiling 计算后下发。
29+struct A5SortSegParam {
30+ int64_t segNum = 0; // 序列被切成多少段
31+ int64_t perSegElements = 0; // 每段元素数(末段可不足)
32+ int64_t totalElements = 0; // 序列总元素数
33+ int64_t expertStart = 0; // 有效 expert_id 的上界:>= 该值者视为被 mask
34+ int64_t sortLenPerSeg = 0; // 每段 proposal 对区的 float 个数(host 按对齐后段长算)
35+};
36+ 
37+class FfnWbA5SegSort {
38+public:
39+ __aicore__ inline FfnWbA5SegSort(){};
40+ 
41+ // flatIdsWs : 归一后的扁平 expert_id(int32)
42+ // pairWs : 段内排序结果的 proposal 对区(float),按段 slice
43+ // cntWs : 每段有效元素数(int32),每段占一个数据块槽位
44+ __aicore__ inline void Init(GM_ADDR flatIdsWs, GM_ADDR pairWs, GM_ADDR cntWs, const A5SortSegParam &param,
45+ const ScheduleContextInfo *ctx, TPipe *pipe)
46+ {
47+ param_ = param;
48+ ctx_ = ctx;
49+ pipe_ = pipe;
50+ // 认领段按**真实向量核编号**:一个 block 下挂多个 AIV 子核时 GetBlockIdx() 相同,
51+ // 只用它认领会让同一段算两遍、另一段没人算。
52+ const int64_t subNum = GetTaskRation() > 0 ? GetTaskRation() : 1;
53+ vecId_ = GetBlockIdx() * subNum + GetSubBlockIdx();
54+ vecNum_ = GetBlockNum() * subNum;
55+ 
56+ flatIdsGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(flatIdsWs), param_.totalElements);
57+ pairGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(pairWs), param_.segNum * param_.sortLenPerSeg);
58+ cntGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(cntWs),
59+ param_.segNum * (ONE_BLK_SIZE / static_cast<int64_t>(sizeof(int32_t))));
60+ 
61+ // UB:输入段(int32 id + int32 原下标) / proposal 对区 / 排序临时区
62+ // 输入区按比较指令的 repeat 粒度对齐：CompareScalar 以 ONE_REPEAT_COMPARE_NUM 为单位处理，
63+ // 传入的元素数向上取整后可能越过段长，缓冲区需按同一粒度留足，避免读写越界。
64+ const int64_t alignSeg = Ceil(param_.perSegElements, ONE_REPEAT_COMPARE_NUM) * ONE_REPEAT_COMPARE_NUM;
65+ pipe_->InitBuffer(inQue_, 1, alignSeg * NUM_TWO * sizeof(int32_t) + ONE_BLK_SIZE);
66+ pipe_->InitBuffer(concatBuf_, param_.sortLenPerSeg * sizeof(float) + ONE_BLK_SIZE);
67+ pipe_->InitBuffer(tmpBuf_, param_.sortLenPerSeg * sizeof(float) + ONE_BLK_SIZE);
68+ pipe_->InitBuffer(sortedBuf_, param_.sortLenPerSeg * sizeof(float) + ONE_BLK_SIZE);
69+ pipe_->InitBuffer(maskBuf_, alignSeg * sizeof(uint32_t) + ONE_BLK_SIZE);
70+ }
71+ 
72+ __aicore__ inline void Process()
73+ {
74+ for (int64_t seg = vecId_; seg < param_.segNum; seg += vecNum_) {
75+ ProcessOneSeg(seg);
76+ }
77+ }
78+ 
79+private:
80+ // 取负 + 判定有效位。取负让降序排序等价于 expert_id 升序;判定阈值同样取负后比较,
81+ // 于是"有效"= 键 > -expertStart。这一段是纯规则的逐元素向量运算,用 MicroAPI 写。
82+ __aicore__ inline void PrepareKeys(const LocalTensor<float> &keys, const LocalTensor<uint32_t> &maskBits,
83+ int64_t count)
84+ {
85+ const uint16_t repeatTimes = static_cast<uint16_t>(Ceil(count, FLOAT_REG_ELEMENTS));
86+ uint32_t remain = static_cast<uint32_t>(count);
87+ __local_mem__ float *keyAddr = reinterpret_cast<__local_mem__ float *>(keys.GetPhyAddr());
88+ const float negOne = -1.0f;
89+ 
90+ __VEC_SCOPE__
91+ {
92+ MicroAPI::MaskReg loopMask;
93+ MicroAPI::RegTensor<float> keyReg;
94+ for (uint16_t i = 0; i < repeatTimes; i++) {
95+ loopMask = MicroAPI::UpdateMask<float>(remain);
96+ MicroAPI::DataCopy(keyReg, keyAddr + i * FLOAT_REG_ELEMENTS);
97+ MicroAPI::Muls(keyReg, keyReg, negOne, loopMask);
98+ MicroAPI::DataCopy(keyAddr + i * FLOAT_REG_ELEMENTS, keyReg, loopMask);
99+ }
100+ }
101+ PipeBarrier<PIPE_V>();
102+ // 有效位:键 > -expertStart(即原 expert_id < expertStart)。比较按 64 元素粒度对齐。
103+ LocalTensor<uint8_t> maskU8 = maskBits.template ReinterpretCast<uint8_t>();
104+ CompareScalar(maskU8, keys, static_cast<float>(-param_.expertStart), CMPMODE::GT,
105+ Ceil(count, ONE_REPEAT_COMPARE_NUM) * ONE_REPEAT_COMPARE_NUM);
106+ PipeBarrier<PIPE_V>();
107+ }
108+ 
109+ __aicore__ inline void ProcessOneSeg(int64_t seg)
110+ {
111+ const int64_t begin = seg * param_.perSegElements;
112+ const int64_t len =
113+ (begin >= param_.totalElements) ? 0 : Min(param_.perSegElements, param_.totalElements - begin);
114+ 
115+ LocalTensor<int32_t> inLocal = inQue_.AllocTensor<int32_t>();
116+ int64_t validCnt = 0;
117+ validCnt = SortOneSegBody(inLocal, seg, begin, len);
118+ 
119+ // 本段有效数写入计数区:跨核可见的交换一律经 UB + DataCopyPad,不用 GM 标量写。
120+ LocalTensor<int32_t> cntLocal = inLocal;
121+ cntLocal.SetValue(0, static_cast<int32_t>(validCnt));
122+ SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
123+ DataCopyExtParams cpCnt{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(int32_t)), 0, 0, 0};
124+ DataCopyPad(cntGm_[seg * (ONE_BLK_SIZE / static_cast<int64_t>(sizeof(int32_t)))], cntLocal, cpCnt);
125+ // 下一段会用 MTE2 重新载入、用向量指令重写同一批缓冲:必须等本段的搬出真正读完。
126+ SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
127+ SetWaitFlag<HardEvent::MTE3_V>(HardEvent::MTE3_V);
128+ inQue_.FreeTensor(inLocal);
129+ }
130+ 
131+ // 段内排序主体:压掉被 mask 的 token,补尾到 Sort32 粒度后 Concat+Sort,结果落 pairGm_。
132+ // 返回本段的有效元素数。
133+ __aicore__ inline int64_t SortOneSegBody(const LocalTensor<int32_t> &inLocal, int64_t seg, int64_t begin,
134+ int64_t len)
135+ {
136+ int64_t validCnt = 0;
137+ if (len > 0) {
138+ DataCopyExtParams cp{static_cast<uint16_t>(1), static_cast<uint32_t>(len * sizeof(int32_t)), 0, 0, 0};
139+ DataCopyPadExtParams<int32_t> pad{false, 0, 0, 0};
140+ DataCopyPad(inLocal, flatIdsGm_[begin], cp, pad);
141+ 
142+ // 原下标:排序的 payload。段内以全局下标编号,归并后即 gather_idx。
143+ const int64_t alignSeg = Ceil(param_.perSegElements, ONE_REPEAT_COMPARE_NUM) * ONE_REPEAT_COMPARE_NUM;
144+ LocalTensor<int32_t> idxLocal = inLocal[alignSeg];
145+ SetWaitFlag<HardEvent::MTE2_V>(HardEvent::MTE2_V);
146+ ArithProgression<int32_t>(idxLocal, static_cast<int32_t>(begin), 1, len);
147+ PipeBarrier<PIPE_V>();
148+ 
149+ LocalTensor<float> keys = inLocal.template ReinterpretCast<float>();
150+ Cast(keys, inLocal, RoundMode::CAST_ROUND, len);
151+ PipeBarrier<PIPE_V>();
152+ 
153+ LocalTensor<uint32_t> maskBits = maskBuf_.Get<uint32_t>();
154+ PrepareKeys(keys, maskBits, len);
155+ 
156+ // 压掉被 mask 的 token:键与下标用同一套有效位压缩,压缩后个数即本段有效数。
157+ uint64_t rsvd = 0;
158+ GatherMaskParams gp;
159+ gp.repeatTimes = 1;
160+ gp.src0BlockStride = 1;
161+ gp.src0RepeatStride = BLOCKS_PER_REPEAT;
162+ gp.src1RepeatStride = 0;
163+ GatherMask(keys, keys, maskBits, true, static_cast<uint32_t>(len), gp, rsvd);
164+ PipeBarrier<PIPE_V>();
165+ LocalTensor<uint32_t> idxU32 = idxLocal.template ReinterpretCast<uint32_t>();
166+ GatherMask(idxU32, idxU32, maskBits, true, static_cast<uint32_t>(len), gp, rsvd);
167+ PipeBarrier<PIPE_V>();
168+ validCnt = static_cast<int64_t>(rsvd);
169+ 
170+ if (validCnt > 0) {
171+ // 尾部补最小值到 Sort32 粒度:补位在降序中落到末尾,不影响有效元素次序。
172+ const int64_t alignCnt = Ceil(validCnt, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
173+ const int64_t tailNum = validCnt % ONE_REPEAT_SORT_NUM;
174+ if (tailNum > 0) {
175+ // 补尾必须从**对齐的起点**发起:向量指令要求 UB 地址按数据块对齐,
176+ // 直接从 keys[validCnt] 开始会触发 "VEC access UB not aligned" 硬件异常。
177+ // 故起点退到本 repeat 的开头,用位掩码只写尾部那几个元素。
178+ uint64_t maskBits = UINT64_MAX << tailNum;
179+ maskBits &= (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
180+ uint64_t dupMask[NUM_TWO] = {maskBits, 0};
181+ Duplicate(keys[validCnt - tailNum], SORT_FILL_VALUE, dupMask, 1, DST_BLK_STRIDE, DST_REP_STRIDE);
182+ PipeBarrier<PIPE_V>();
183+ }
184+ // Concat/Sort 的各 tensor 不允许地址重叠:concat 结果、临时区、排序结果各占一块。
185+ LocalTensor<float> concatLocal = concatBuf_.Get<float>();
186+ LocalTensor<float> tmpLocal = tmpBuf_.Get<float>();
187+ LocalTensor<float> sortedLocal = sortedBuf_.Get<float>();
188+ Concat(concatLocal, keys, tmpLocal, alignCnt / ONE_REPEAT_SORT_NUM);
189+ PipeBarrier<PIPE_V>();
190+ Sort<float, true>(sortedLocal, concatLocal, idxU32, tmpLocal, alignCnt / ONE_REPEAT_SORT_NUM);
191+ PipeBarrier<PIPE_V>();
192+ 
193+ SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
194+ DataCopyExtParams cpOut{static_cast<uint16_t>(1),
195+ static_cast<uint32_t>(GetSortLen<float>(alignCnt) * sizeof(float)), 0, 0, 0};
196+ DataCopyPad(pairGm_[seg * param_.sortLenPerSeg], sortedLocal, cpOut);
197+ SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
198+ }
199+ }
200+ return validCnt;
201+ }
202+ 
203+private:
204+ // 一个向量寄存器容纳的 fp32 个数,以及一次 repeat 覆盖的数据块数/比较元素数。
205+ static constexpr int64_t FLOAT_REG_ELEMENTS = 64;
206+ static constexpr int64_t BLOCKS_PER_REPEAT = 8;
207+ static constexpr int64_t ONE_REPEAT_COMPARE_NUM = 64;
208+ // Duplicate 的目的地块内步长与 repeat 间步长(按数据块计),取同仓 arch35 算子同款配置。
209+ static constexpr int64_t DST_BLK_STRIDE = 1;
210+ static constexpr int64_t DST_REP_STRIDE = 8;
211+ 
212+ A5SortSegParam param_;
213+ const ScheduleContextInfo *ctx_ = nullptr;
214+ TPipe *pipe_ = nullptr;
215+ 
216+ GlobalTensor<int32_t> flatIdsGm_;
217+ GlobalTensor<float> pairGm_;
218+ GlobalTensor<int32_t> cntGm_;
219+ 
220+ TQue<QuePosition::VECIN, 1> inQue_;
221+ TBuf<TPosition::VECCALC> concatBuf_;
222+ TBuf<TPosition::VECCALC> tmpBuf_;
223+ TBuf<TPosition::VECCALC> sortedBuf_;
224+ TBuf<TPosition::VECCALC> maskBuf_;
225+ 
226+ int64_t vecId_ = 0;
227+ int64_t vecNum_ = 0;
228+};
229+ 
230+// ===================== 段间归并 =====================
231+// 归并的切分参数,由 host tiling 计算下发。
232+struct A5MergeParam {
233+ int64_t segNum = 0; // 段数(= 第 0 轮的路数)
234+ int64_t sortLenPerSeg = 0; // 每段 slot 的 float 个数
235+ int64_t oneLoopMaxElements = 0; // 单路单次载入 UB 的最大元素数
236+ int64_t rounds = 0; // 归并轮数 = ceil(log4(segNum))
237+};
238+ 
239+class FfnWbA5MrgSort {
240+public:
241+ __aicore__ inline FfnWbA5MrgSort(){};
242+ 
243+ // wsA/wsB:两块等大的 proposal 对工作区(乒乓);cntWs:各段有效元素数
244+ __aicore__ inline void Init(GM_ADDR wsA, GM_ADDR wsB, GM_ADDR cntWs, const A5MergeParam &param, TPipe *pipe)
245+ {
246+ param_ = param;
247+ pipe_ = pipe;
248+ const int64_t subNum = GetTaskRation() > 0 ? GetTaskRation() : 1;
249+ vecId_ = GetBlockIdx() * subNum + GetSubBlockIdx();
250+ vecNum_ = GetBlockNum() * subNum;
251+ 
252+ const int64_t wsFloats = param_.segNum * param_.sortLenPerSeg;
253+ wsGm_[0].SetGlobalBuffer(reinterpret_cast<__gm__ float *>(wsA), wsFloats);
254+ wsGm_[1].SetGlobalBuffer(reinterpret_cast<__gm__ float *>(wsB), wsFloats);
255+ cntGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(cntWs),
256+ param_.segNum * (ONE_BLK_SIZE / static_cast<int64_t>(sizeof(int32_t))));
257+ 
258+ // UB:MRG_LIST_NUM 路输入 + 1 路输出,均按 proposal 对宽度分配
259+ const int64_t loopFloats = GetSortLen<float>(param_.oneLoopMaxElements);
260+ pipe_->InitBuffer(inQue_, 1, loopFloats * MRG_LIST_NUM * sizeof(float) + ONE_BLK_SIZE);
261+ pipe_->InitBuffer(outQue_, 1, loopFloats * MRG_LIST_NUM * sizeof(float) + ONE_BLK_SIZE);
262+ pipe_->InitBuffer(cntBuf_,
263+ Align(param_.segNum * (ONE_BLK_SIZE / sizeof(int32_t)), sizeof(int32_t)) * sizeof(int32_t) +
264+ ONE_BLK_SIZE);
265+ }
266+ 
267+ // 跑完全部归并轮次;返回最终结果所在的工作区序号(0=wsA,1=wsB)。
268+ // 每轮内各组由不同向量核认领,轮与轮之间由调用方 SyncAll。
269+ __aicore__ inline int64_t ProcessRound(int64_t round, int64_t srcIdx)
270+ {
271+ LoadSegCounts();
272+ const int64_t groupStride = Pow4(round + 1); // 本轮一组覆盖的段数
273+ const int64_t listStride = Pow4(round); // 组内相邻两路相隔的段数
274+ const int64_t groupNum = Ceil(param_.segNum, groupStride);
275+ 
276+ for (int64_t g = vecId_; g < groupNum; g += vecNum_) {
277+ MergeOneGroup(g, groupStride, listStride, srcIdx);
278+ }
279+ return 1 - srcIdx;
280+ }
281+ 
282+ // 最终一路的元素总数 = 各段有效数之和
283+ __aicore__ inline int64_t TotalValid()
284+ {
285+ LoadSegCounts();
286+ int64_t total = 0;
287+ for (int64_t s = 0; s < param_.segNum; s++) {
288+ total += SegCount(s);
289+ }
290+ return total;
291+ }
292+ 
293+private:
294+ __aicore__ inline int64_t Pow4(int64_t e)
295+ {
296+ int64_t v = 1;
297+ for (int64_t i = 0; i < e; i++) {
298+ v *= MRG_LIST_NUM;
299+ }
300+ return v;
301+ }
302+ 
303+ __aicore__ inline void LoadSegCounts()
304+ {
305+ if (cntLoaded_) {
306+ return;
307+ }
308+ cntLocal_ = cntBuf_.Get<int32_t>();
309+ const int64_t words = param_.segNum * (ONE_BLK_SIZE / static_cast<int64_t>(sizeof(int32_t)));
310+ DataCopyExtParams cp{static_cast<uint16_t>(1), static_cast<uint32_t>(words * sizeof(int32_t)), 0, 0, 0};
311+ DataCopyPadExtParams<int32_t> pad{false, 0, 0, 0};
312+ DataCopyPad(cntLocal_, cntGm_, cp, pad);
313+ SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
314+ cntLoaded_ = true;
315+ }
316+ 
317+ __aicore__ inline int64_t SegCount(int64_t seg)
318+ {
319+ if (seg >= param_.segNum) {
320+ return 0;
321+ }
322+ return static_cast<int64_t>(cntLocal_.GetValue(seg * (ONE_BLK_SIZE / sizeof(int32_t))));
323+ }
324+ 
325+ // 第 r 轮中,以 seg 为首段、跨度 listStride 的那一路的元素数 = 其覆盖段的有效数之和
326+ __aicore__ inline int64_t ListLength(int64_t firstSeg, int64_t listStride)
327+ {
328+ int64_t sum = 0;
329+ for (int64_t s = firstSeg; s < firstSeg + listStride && s < param_.segNum; s++) {
330+ sum += SegCount(s);
331+ }
332+ return sum;
333+ }
334+ 
335+ __aicore__ inline void MergeOneGroup(int64_t group, int64_t groupStride, int64_t listStride, int64_t srcIdx)
336+ {
337+ const int64_t baseSeg = group * groupStride;
338+ int64_t offsets[MRG_LIST_NUM];
339+ int64_t remains[MRG_LIST_NUM];
340+ int64_t listNum = 0;
341+ for (int64_t i = 0; i < MRG_LIST_NUM; i++) {
342+ const int64_t firstSeg = baseSeg + i * listStride;
343+ if (firstSeg >= param_.segNum) {
344+ break;
345+ }
346+ const int64_t len = ListLength(firstSeg, listStride);
347+ offsets[listNum] = firstSeg * param_.sortLenPerSeg;
348+ remains[listNum] = len;
349+ listNum++;
350+ }
351+ if (listNum == 0) {
352+ return;
353+ }
354+ 
355+ LocalTensor<float> inLocal = inQue_.AllocTensor<float>();
356+ LocalTensor<float> outLocal = outQue_.AllocTensor<float>();
357+ const int64_t loopFloats = GetSortLen<float>(param_.oneLoopMaxElements);
358+ int64_t outOffset = baseSeg * param_.sortLenPerSeg;
359+ 
360+ int64_t allRemain = 0;
361+ for (int64_t i = 0; i < listNum; i++) {
362+ allRemain += remains[i];
363+ }
364+ 
365+ MergeStream(inLocal, outLocal, listNum, loopFloats, outOffset, allRemain, offsets, remains, srcIdx);
366+ 
367+ inQue_.FreeTensor(inLocal);
368+ outQue_.FreeTensor(outLocal);
369+ }
370+ 
371+ // 流式归并主循环:每轮从各路取一段填满 UB,MrgSort 后整段写出,直到各路耗尽。
372+ // 从各路各取一块填进 UB 输入区(已耗尽的路跳过),记下每块元素数与其对应的原路号。
373+ // 返回本轮真正参与归并的路数。
374+ __aicore__ inline int64_t FillLists(const LocalTensor<float> &inLocal, int64_t listNum, int64_t loopFloats,
375+ int64_t srcIdx, const int64_t (&offsets)[MRG_LIST_NUM],
376+ const int64_t (&remains)[MRG_LIST_NUM],
377+ LocalTensor<float> (&lists)[MRG_LIST_NUM], uint16_t (&counts)[MRG_LIST_NUM],
378+ int64_t (&liveMap)[MRG_LIST_NUM])
379+ {
380+ int64_t liveNum = 0;
381+ for (int64_t i = 0; i < listNum; i++) {
382+ const int64_t take = Min(param_.oneLoopMaxElements, remains[i]);
383+ if (take <= 0) {
384+ continue;
385+ }
386+ DataCopyExtParams cp{static_cast<uint16_t>(1),
387+ static_cast<uint32_t>(GetSortLen<float>(take) * sizeof(float)), 0, 0, 0};
388+ DataCopyPadExtParams<float> pad{false, 0, 0, 0};
389+ DataCopyPad(inLocal[liveNum * loopFloats], wsGm_[srcIdx][offsets[i]], cp, pad);
390+ lists[liveNum] = inLocal[liveNum * loopFloats];
391+ counts[liveNum] = static_cast<uint16_t>(take);
392+ liveMap[liveNum] = i;
393+ liveNum++;
394+ }
395+ return liveNum;
396+ }
397+ 
398+ __aicore__ inline void MergeStream(const LocalTensor<float> &inLocal, const LocalTensor<float> &outLocal,
399+ int64_t listNum, int64_t loopFloats, int64_t &outOffset, int64_t &allRemain,
400+ int64_t (&offsets)[MRG_LIST_NUM], int64_t (&remains)[MRG_LIST_NUM],
401+ int64_t srcIdx)
402+ {
403+ while (allRemain > 0) {
404+ // 载入各路当前块
405+ uint16_t counts[MRG_LIST_NUM] = {0, 0, 0, 0};
406+ LocalTensor<float> lists[MRG_LIST_NUM];
407+ int64_t liveMap[MRG_LIST_NUM];
408+ const int64_t liveNum =
409+ FillLists(inLocal, listNum, loopFloats, srcIdx, offsets, remains, lists, counts, liveMap);
410+ if (liveNum == 0) {
411+ break;
412+ }
413+ SetWaitFlag<HardEvent::MTE2_V>(HardEvent::MTE2_V);
414+ 
415+ uint32_t sortedNums[MRG_LIST_NUM] = {0, 0, 0, 0};
416+ if (liveNum == 1) {
417+ // 只剩一路:直接搬,不必归并
418+ DataCopy(outLocal, lists[0], Align(GetSortLen<float>(counts[0]), sizeof(float)));
419+ sortedNums[0] = counts[0];
420+ } else {
421+ MrgSortSrcList srcList =
422+ MrgSortSrcList(lists[0], lists[liveNum > 1 ? 1 : 0], lists[liveNum > NUM_TWO ? NUM_TWO : 0],
423+ lists[liveNum > NUM_THREE ? NUM_THREE : 0]);
424+ const uint16_t validBit = static_cast<uint16_t>((1U << liveNum) - 1U);
425+ MrgSort<float, true>(outLocal, srcList, counts, sortedNums, validBit, 1);
426+ }
427+ PipeBarrier<PIPE_V>();
428+ 
429+ // 推进各路偏移与剩余
430+ int64_t produced = 0;
431+ for (int64_t j = 0; j < liveNum; j++) {
432+ const int64_t consumed = static_cast<int64_t>(sortedNums[j]);
433+ const int64_t i = liveMap[j];
434+ offsets[i] += GetSortLen<float>(consumed);
435+ remains[i] -= consumed;
436+ allRemain -= consumed;
437+ produced += consumed;
438+ }
439+ if (produced <= 0) {
440+ break; // 保护:不应发生,避免死循环
441+ }
442+ 
443+ SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
444+ DataCopyExtParams cpOut{static_cast<uint16_t>(1),
445+ static_cast<uint32_t>(GetSortLen<float>(produced) * sizeof(float)), 0, 0, 0};
446+ DataCopyPad(wsGm_[1 - srcIdx][outOffset], outLocal, cpOut);
447+ // 下一轮要用 MTE2 载入输入、用 MrgSort(向量)重写输出缓冲,两者都要等本次搬出读完。
448+ SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
449+ SetWaitFlag<HardEvent::MTE3_V>(HardEvent::MTE3_V);
450+ outOffset += GetSortLen<float>(produced);
451+ }
452+ }
453+ 
454+private:
455+ A5MergeParam param_;
456+ TPipe *pipe_ = nullptr;
457+ 
458+ GlobalTensor<float> wsGm_[NUM_TWO];
459+ GlobalTensor<int32_t> cntGm_;
460+ 
461+ TQue<QuePosition::VECIN, 1> inQue_;
462+ TQue<QuePosition::VECOUT, 1> outQue_;
463+ TBuf<TPosition::VECCALC> cntBuf_;
464+ LocalTensor<int32_t> cntLocal_;
465+ bool cntLoaded_ = false;
466+ 
467+ int64_t vecId_ = 0;
468+ int64_t vecNum_ = 0;
469+};
470+ 
471+// ===================== 归并收尾:拆出 id 与 idx =====================
472+struct A5ExtractParam {
473+ int64_t totalValid = 0; // 最终有序序列的元素数
474+ int64_t perLoopElements = 0; // 单次载入 UB 的元素数(host 保证为 Sort32 粒度的整数倍)
475+};
476+ 
477+class FfnWbA5Extract {
478+public:
479+ __aicore__ inline FfnWbA5Extract(){};
480+ 
481+ // pairWs : 最终一路 proposal 对
482+ // idsWs : 输出,有序 expert_id(int32)
483+ // idxWs : 输出,gather_idx(int32,指向原始扁平序列的下标)
484+ __aicore__ inline void Init(GM_ADDR pairWs, GM_ADDR idsWs, GM_ADDR idxWs, const A5ExtractParam &param, TPipe *pipe)
485+ {
486+ param_ = param;
487+ pipe_ = pipe;
488+ const int64_t subNum = GetTaskRation() > 0 ? GetTaskRation() : 1;
489+ vecId_ = GetBlockIdx() * subNum + GetSubBlockIdx();
490+ vecNum_ = GetBlockNum() * subNum;
491+ 
492+ pairGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(pairWs), GetSortLen<float>(param_.totalValid));
493+ idsGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(idsWs), param_.totalValid);
494+ idxGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(idxWs), param_.totalValid);
495+ 
496+ // Extract 按 ONE_REPEAT_SORT_NUM 成批写出:落盘量是**对齐后的长度**,不是有效长度。
497+ // 按有效长度分配会溢出到相邻缓冲(实测把索引缓冲的开头冲成 0),故三块都按对齐长度给。
498+ const int64_t alignLoop = Ceil(param_.perLoopElements, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
499+ pipe_->InitBuffer(pairQue_, 1, GetSortLen<float>(alignLoop) * sizeof(float) + ONE_BLK_SIZE);
500+ pipe_->InitBuffer(idsBuf_, alignLoop * sizeof(int32_t) + ONE_BLK_SIZE);
501+ pipe_->InitBuffer(idxBuf_, alignLoop * sizeof(int32_t) + ONE_BLK_SIZE);
502+ }
503+ 
504+ __aicore__ inline void Process()
505+ {
506+ const int64_t loops = Ceil(param_.totalValid, param_.perLoopElements);
507+ for (int64_t l = vecId_; l < loops; l += vecNum_) {
508+ const int64_t begin = l * param_.perLoopElements;
509+ const int64_t len = Min(param_.perLoopElements, param_.totalValid - begin);
510+ if (len <= 0) {
511+ continue;
512+ }
513+ ProcessOneChunk(begin, len);
514+ }
515+ }
516+ 
517+private:
518+ __aicore__ inline void ProcessOneChunk(int64_t begin, int64_t len)
519+ {
520+ // Extract 按 Sort32 粒度成批工作,故按对齐后的长度拆包,超出有效长度的部分不写出。
521+ const int64_t alignLen = Ceil(len, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
522+ 
523+ LocalTensor<float> pairLocal = pairQue_.AllocTensor<float>();
524+ DataCopyExtParams cpIn{static_cast<uint16_t>(1), static_cast<uint32_t>(GetSortLen<float>(len) * sizeof(float)),
525+ 0, 0, 0};
526+ DataCopyPadExtParams<float> pad{false, 0, 0, 0};
527+ DataCopyPad(pairLocal, pairGm_[GetSortLen<float>(begin)], cpIn, pad);
528+ SetWaitFlag<HardEvent::MTE2_V>(HardEvent::MTE2_V);
529+ 
530+ LocalTensor<int32_t> idsLocal = idsBuf_.Get<int32_t>();
531+ LocalTensor<int32_t> idxLocal = idxBuf_.Get<int32_t>();
532+ LocalTensor<float> keysLocal = idsLocal.template ReinterpretCast<float>();
533+ LocalTensor<uint32_t> idxU32 = idxLocal.template ReinterpretCast<uint32_t>();
534+ 
535+ Extract(keysLocal, idxU32, pairLocal, static_cast<int32_t>(alignLen / ONE_REPEAT_SORT_NUM));
536+ PipeBarrier<PIPE_V>();
537+ // 还原排序前取的负号,再转回 int32
538+ Muls(keysLocal, keysLocal, static_cast<float>(-1), len);
539+ PipeBarrier<PIPE_V>();
540+ Cast(idsLocal, keysLocal, RoundMode::CAST_ROUND, len);
541+ PipeBarrier<PIPE_V>();
542+ 
543+ SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
544+ DataCopyExtParams cpOut{static_cast<uint16_t>(1), static_cast<uint32_t>(len * sizeof(int32_t)), 0, 0, 0};
545+ DataCopyPad(idsGm_[begin], idsLocal, cpOut);
546+ DataCopyPad(idxGm_[begin], idxLocal, cpOut);
547+ // 下一块会用 Extract(向量)重写这两块缓冲:必须等本块搬出读完。
548+ SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
549+ SetWaitFlag<HardEvent::MTE3_V>(HardEvent::MTE3_V);
550+ pairQue_.FreeTensor(pairLocal);
551+ }
552+ 
553+private:
554+ A5ExtractParam param_;
555+ TPipe *pipe_ = nullptr;
556+ 
557+ GlobalTensor<float> pairGm_;
558+ GlobalTensor<int32_t> idsGm_;
559+ GlobalTensor<int32_t> idxGm_;
560+ 
561+ TQue<QuePosition::VECIN, 1> pairQue_;
562+ TBuf<TPosition::VECCALC> idsBuf_;
563+ TBuf<TPosition::VECCALC> idxBuf_;
564+ 
565+ int64_t vecId_ = 0;
566+ int64_t vecNum_ = 0;
567+};
568+} // namespace FfnWbBatchingArch35
569+#endif // OP_KERNEL_ARCH35_FFN_WB_A5_SORT_H
Dffn/ffn_worker_batching/op_kernel/arch35/ffn_wb_arch35_reuse.h+0-45
@@ -1,45 +0,0 @@
1-/**
2- * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
3- *
4- * Licensed under the Apache License, Version 2.0 (the "License");
5- * you may not use this file except in compliance with the License.
6- * You may obtain a copy of the License at
7- *
8- * http://www.apache.org/licenses/LICENSE-2.0
9- *
10- * Unless required by applicable law or agreed to in writing, software
11- * distributed under the License is distributed on an "AS IS" BASIS,
12- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13- * See the License for the specific language governing permissions and
14- * limitations under the License.
15- */
16- 
17-/*!
18- * \file ffn_wb_arch35_reuse.h
19- * \brief arch35(A5) 复用清单:以下 kernel 头与 A2 逻辑完全一致,直接复用 A2 实现,经命名空间桥接进
20- * FfnWbBatchingArch35(A2/A3 零改动)。arch35 仅保留有真实代次差异的 *_arch35.h 独立实现。
21- */
22-#ifndef OP_KERNEL_ARCH35_FFN_WB_ARCH35_REUSE_H
23-#define OP_KERNEL_ARCH35_FFN_WB_ARCH35_REUSE_H
24-#include "../ffn_wb_common.h"
25-#include "../ffn_wb_sort_base.h"
26-#include "../ffn_wb_sort_mrgsort.h"
27-#include "../ffn_wb_sort_mrgsort_out.h"
28-#include "../ffn_wb_sort_one_core.h"
29-#include "../ffn_wb_get_schedule_context.h"
30-#include "../ffn_wb_gather_out_all.h"
31-#include "../ffn_wb_scan_token_info.h"
32-#include "../ffn_wb_scan_sort_one_core.h"
33-#include "../ffn_wb_scan_get_valid_experts.h"
34-namespace FfnWbBatchingArch35 {
35-using namespace FfnWbBatching;
36-// op 辅助函数与 AscendC 同名,using-声明取优先解析,消除二义。
37-using FfnWbBatching::Align;
38-using FfnWbBatching::Ceil;
39-using FfnWbBatching::CeilDiv;
40-using FfnWbBatching::Max;
41-using FfnWbBatching::Min;
42-using FfnWbBatching::PowerOfFourCeil;
43-using FfnWbBatching::SetWaitFlag;
44-} // namespace FfnWbBatchingArch35
45-#endif // OP_KERNEL_ARCH35_FFN_WB_ARCH35_REUSE_H
Dffn/ffn_worker_batching/op_kernel/arch35/ffn_wb_group_listing_arch35.h+0-332
@@ -1,332 +0,0 @@
1-/**
2- * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
3- *
4- * Licensed under the Apache License, Version 2.0 (the "License");
5- * you may not use this file except in compliance with the License.
6- * You may obtain a copy of the License at
7- *
8- * http://www.apache.org/licenses/LICENSE-2.0
9- *
10- * Unless required by applicable law or agreed to in writing, software
11- * distributed under the License is distributed on an "AS IS" BASIS,
12- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13- * See the License for the specific language governing permissions and
14- * limitations under the License.
15- */
16- 
17-/*!
18- * \file ffn_wb_group_listing_arch35.h
19- * \brief
20- */
21- 
22-#ifndef OP_KERNEL_FFN_WB_GROUP_LISTING_H
23-#define OP_KERNEL_FFN_WB_GROUP_LISTING_H
24-#include "ffn_wb_arch35_reuse.h"
25- 
26-#include "kernel_operator.h"
27- 
28-namespace FfnWbBatchingArch35 {
29-using namespace AscendC;
30- 
31-class FfnWbGroupListingArch35 {
32-public:
33- __aicore__ inline FfnWbGroupListingArch35(){};
34- __aicore__ inline void Init(GM_ADDR x, GM_ADDR groupList, GM_ADDR groupListTmp,
35- ScheduleContextInfo *scheduleContext, TPipe *tPipe, int64_t groupListingDealFlag);
36- __aicore__ inline void Process(int64_t groupListingDealFlag);
37- __aicore__ inline void ProcessExpertCount();
38- 
39-private:
40- __aicore__ inline void CopyIn(int64_t loop, int64_t curLoopElements);
41- __aicore__ inline void Compute(int64_t curLoopElements);
42- __aicore__ inline void CopyOut();
43- 
44- __aicore__ inline void expertCountCopyIn();
45- __aicore__ inline void expertCountCompute();
46- __aicore__ inline void expertCountCopyOut();
47- __aicore__ inline void CopyInOneCore(int64_t inputNum);
48- __aicore__ inline void ComputeOneCore(int64_t inputNum);
49- 
50-private:
51- TPipe *pipe_;
52- 
53- GlobalTensor<int32_t> expandedExpertIdsGm; // 排序后的专家索引（全局内存）
54- GlobalTensor<int64_t> groupListGm; // 最终专家Token计数（int64）
55- GlobalTensor<int32_t> expertCountTempGm; // 临时专家计数（int32）
56- LocalTensor<int64_t> groupListOutLocal;
57- 
58- TQue<QuePosition::VECIN, 1> sortedExpertIdxInQueue; // 专家索引输入队列
59- TQue<QuePosition::VECOUT, 1> expertIdxCountOutQueue; // 最终计数输出队列
60- 
61- int64_t blockIdx;
62- int64_t perLoopRows;
63- int64_t coreNum = 0;
64- int64_t expertNum = 0;
65- 
66- int64_t tokenCount = 0;
67- int64_t curExpertIdOffset = 0;
68- int32_t lastExpertId = -1;
69- int32_t firstExpertId = -1;
70- 
71- int64_t curCoreElements = 0;
72- int64_t curcoreLoopsNum = 0;
73- int64_t curCorePerLoopElements = 0;
74- int64_t curCoreLastLoopElements = 0;
75- 
76- int64_t perCoreElements = 0;
77- int64_t lastCoreElements = 0;
78- int64_t perCoreLoopsNum = 0;
79- int64_t lastCoreLoopsNum = 0;
80- int64_t perCorePerLoopElements = 0;
81- int64_t perCoreLastLoopElements = 0;
82- int64_t lastCorePerLoopElements = 0;
83- int64_t lastCoreLastLoopElements = 0;
84- int64_t validGatherIdxLength = 0;
85- int64_t actualExpertTotalNum = 0;
86-};
87- 
88-__aicore__ inline void FfnWbGroupListingArch35::Init(GM_ADDR x, GM_ADDR groupList, GM_ADDR groupListTmp,
89- ScheduleContextInfo *scheduleContext, TPipe *tPipe,
90- int64_t groupListingDealFlag)
91-{
92- pipe_ = tPipe;
93- blockIdx = GetBlockIdx();
94- validGatherIdxLength = scheduleContext->validGatherIdxLength;
95- 
96- coreNum = scheduleContext->coreNum;
97- perLoopRows = 8192L;
98- 
99- // 计算每个核处理数据
100- perCoreElements = validGatherIdxLength / GROUP_LISTING_MULTI_AIV_NUM;
101- lastCoreElements = perCoreElements + (validGatherIdxLength % GROUP_LISTING_MULTI_AIV_NUM);
102- 
103- // 计算每个核循环次数
104- perCoreLoopsNum = (perCoreElements + perLoopRows - 1) / perLoopRows;
105- perCorePerLoopElements = perLoopRows;
106- perCoreLastLoopElements = (perCoreElements % perLoopRows == 0) ? perLoopRows : (perCoreElements % perLoopRows);
107- lastCoreLoopsNum = (lastCoreElements + perLoopRows - 1) / perLoopRows;
108- lastCorePerLoopElements = perLoopRows;
109- lastCoreLastLoopElements = (lastCoreElements % perLoopRows == 0) ? perLoopRows : (lastCoreElements % perLoopRows);
110- 
111- const int64_t firstCoreIdx = coreNum - GROUP_LISTING_MULTI_AIV_NUM;
112- if ((blockIdx >= firstCoreIdx) && (blockIdx < coreNum - 1)) {
113- curCoreElements = perCoreElements;
114- curcoreLoopsNum = perCoreLoopsNum;
115- curCorePerLoopElements = perCorePerLoopElements;
116- curCoreLastLoopElements = perCoreLastLoopElements;
117- } else if (blockIdx == coreNum - 1) {
118- curCoreElements = lastCoreElements;
119- curcoreLoopsNum = lastCoreLoopsNum;
120- curCorePerLoopElements = lastCorePerLoopElements;
121- curCoreLastLoopElements = lastCoreLastLoopElements;
122- }
123- 
124- expertNum = scheduleContext->expertNum;
125- 
126- // 输入global 排序好的ExpertIds
127- if (groupListingDealFlag == 1) {
128- expandedExpertIdsGm.SetGlobalBuffer((__gm__ int32_t *)x + (blockIdx - firstCoreIdx) * perCoreElements,
129- curCoreElements);
130- } else {
131- expandedExpertIdsGm.SetGlobalBuffer((__gm__ int32_t *)x, validGatherIdxLength);
132- }
133- // 输出global group_list数据[id * count] int_64
134- groupListGm.SetGlobalBuffer((__gm__ int64_t *)groupList, expertNum * NUM_TWO);
135- // 输出global 存放原子加group_list数据[count] int_32
136- expertCountTempGm.SetGlobalBuffer((__gm__ int32_t *)groupListTmp, expertNum);
137- 
138- int64_t cntSize = (expertNum * NUM_TWO * sizeof(int64_t) + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES;
139- pipe_->InitBuffer(sortedExpertIdxInQueue, 1, perLoopRows * sizeof(int32_t));
140- pipe_->InitBuffer(expertIdxCountOutQueue, 1, cntSize);
141-}
142- 
143-__aicore__ inline void FfnWbGroupListingArch35::Process(int64_t groupListingDealFlag)
144-{
145- if ((blockIdx >= coreNum - GROUP_LISTING_MULTI_AIV_NUM) && (blockIdx < coreNum) && (groupListingDealFlag == 1)) {
146- for (int64_t i = 0; i < curcoreLoopsNum; i++) {
147- int64_t perLoopElements = (i == (curcoreLoopsNum - 1)) ? curCoreLastLoopElements : curCorePerLoopElements;
148- CopyIn(i, perLoopElements);
149- Compute(perLoopElements);
150- CopyOut();
151- }
152- }
153- 
154- if ((blockIdx == coreNum - 1) && (groupListingDealFlag == 0)) {
155- groupListOutLocal = expertIdxCountOutQueue.AllocTensor<int64_t>();
156- curExpertIdOffset = 0;
157- CopyInOneCore(validGatherIdxLength);
158- ComputeOneCore(validGatherIdxLength);
159- 
160- // 更新最后一个expertId和tokenCnt，并且需要在groupList后面添加一个[0， 0]
161- // k=0(validGatherIdxLength==0)时跳过 lastExpertId flush：此时 ComputeOneCore 读到的是
162- // 0字节搬运后未初始化的 UB 脏值(见 issue k-zero-dirty-expertid)，flush 会写出伪专家条目；
163- // 守卫后仅由下方 [0,0] 终止符块产出 group_list[0]=[0,0]，对齐 spec k_zero 契约。
164- if (validGatherIdxLength > 0 && curExpertIdOffset < expertNum) {
165- groupListOutLocal.SetValue(curExpertIdOffset * NUM_TWO, lastExpertId);
166- groupListOutLocal.SetValue(curExpertIdOffset * NUM_TWO + 1, tokenCount);
167- curExpertIdOffset += 1;
168- }
169- // group list未满的情况下，最后补充[0， 0]
170- if (curExpertIdOffset < expertNum) {
171- groupListOutLocal.SetValue(curExpertIdOffset * NUM_TWO, 0);
172- groupListOutLocal.SetValue(curExpertIdOffset * NUM_TWO + 1, 0);
173- curExpertIdOffset += 1;
174- }
175- 
176- DataCopyExtParams copyParams{static_cast<uint16_t>(1),
177- static_cast<uint32_t>(curExpertIdOffset * NUM_TWO * sizeof(int64_t)), 0, 0, 0};
178- SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
179- DataCopyPad(groupListGm, groupListOutLocal, copyParams);
180- 
181- expertIdxCountOutQueue.FreeTensor(groupListOutLocal);
182- }
183-}
184- 
185-__aicore__ inline void FfnWbGroupListingArch35::ProcessExpertCount()
186-{
187- if (blockIdx == coreNum - 1) {
188- expertCountCopyIn();
189- expertCountCompute();
190- expertCountCopyOut();
191- }
192-}
193- 
194-__aicore__ inline void FfnWbGroupListingArch35::CopyIn(int64_t loop, int64_t curLoopElements)
195-{
196- LocalTensor<int32_t> sortedExpertIdxInLocal = sortedExpertIdxInQueue.AllocTensor<int32_t>();
197- DataCopyExtParams dataCopyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(curLoopElements * sizeof(int32_t)),
198- 0, 0, 0};
199- DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0};
200- int64_t sortedexpertIdxOffset = loop * curCorePerLoopElements;
201- DataCopyPad(sortedExpertIdxInLocal, expandedExpertIdsGm[sortedexpertIdxOffset], dataCopyParams, dataCopyPadParams);
202- 
203- sortedExpertIdxInQueue.EnQue(sortedExpertIdxInLocal);
204-}
205- 
206-__aicore__ inline void FfnWbGroupListingArch35::Compute(int64_t curLoopElements)
207-{
208- LocalTensor<int32_t> sortedExpertIdxInLocal = sortedExpertIdxInQueue.DeQue<int32_t>();
209- LocalTensor<int32_t> expertCountOutLocal = expertIdxCountOutQueue.AllocTensor<int32_t>();
210- Duplicate(expertCountOutLocal.ReinterpretCast<int32_t>(), static_cast<int32_t>(0), static_cast<int32_t>(expertNum));
211- event_t eventIDVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
212- SetFlag<HardEvent::V_S>(eventIDVToS);
213- WaitFlag<HardEvent::V_S>(eventIDVToS);
214- lastExpertId = sortedExpertIdxInLocal.GetValue(0);
215- tokenCount = 1;
216- for (int64_t i = 1; i < curLoopElements; i++) {
217- int32_t curExpertId = sortedExpertIdxInLocal.GetValue(i);
218- tokenCount++;
219- if (curExpertId > lastExpertId) {
220- // 先更新expertId切换前的数据
221- expertCountOutLocal.SetValue(lastExpertId, tokenCount - 1);
222- tokenCount = 1;
223- lastExpertId = curExpertId;
224- }
225- }
226- expertCountOutLocal.SetValue(lastExpertId, tokenCount);
227- 
228- event_t eventIDSToMTE3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3));
229- SetFlag<HardEvent::S_MTE3>(eventIDSToMTE3);
230- WaitFlag<HardEvent::S_MTE3>(eventIDSToMTE3);
231- expertIdxCountOutQueue.EnQue<int32_t>(expertCountOutLocal);
232- sortedExpertIdxInQueue.FreeTensor(sortedExpertIdxInLocal);
233-}
234- 
235-__aicore__ inline void FfnWbGroupListingArch35::CopyOut()
236-{
237- LocalTensor<int32_t> expertCountOutLocal = expertIdxCountOutQueue.DeQue<int32_t>();
238- 
239- DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>((expertNum) * sizeof(int32_t)), 0, 0,
240- 0};
241- SetAtomicAdd<int32_t>();
242- DataCopyPad(expertCountTempGm, expertCountOutLocal, copyParams);
243- SetAtomicNone();
244- 
245- expertIdxCountOutQueue.FreeTensor(expertCountOutLocal);
246-}
247- 
248-__aicore__ inline void FfnWbGroupListingArch35::expertCountCopyIn()
249-{
250- LocalTensor<int32_t> expertCountTempInLocal = sortedExpertIdxInQueue.AllocTensor<int32_t>();
251- 
252- DataCopyExtParams dataCopyParams{static_cast<uint16_t>(1), static_cast<uint32_t>((expertNum) * sizeof(int32_t)), 0,
253- 0, 0};
254- DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0};
255- DataCopyPad(expertCountTempInLocal, expertCountTempGm, dataCopyParams, dataCopyPadParams);
256- 
257- sortedExpertIdxInQueue.EnQue(expertCountTempInLocal);
258-}
259- 
260-__aicore__ inline void FfnWbGroupListingArch35::expertCountCompute()
261-{
262- LocalTensor<int32_t> expertCountTempInLocal = sortedExpertIdxInQueue.DeQue<int32_t>();
263- 
264- LocalTensor<int64_t> expertCountOutLocal = expertIdxCountOutQueue.AllocTensor<int64_t>();
265- event_t eventIDMte2ToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
266- SetFlag<HardEvent::MTE2_S>(eventIDMte2ToS);
267- WaitFlag<HardEvent::MTE2_S>(eventIDMte2ToS);
268- for (int64_t i = 0; i < expertNum; i++) {
269- int64_t expertCount = static_cast<int64_t>(expertCountTempInLocal.GetValue(i));
270- if (expertCount > 0) {
271- actualExpertTotalNum++;
272- expertCountOutLocal.SetValue((actualExpertTotalNum - 1) * NUM_TWO, i);
273- expertCountOutLocal.SetValue((actualExpertTotalNum - 1) * NUM_TWO + 1, expertCount);
274- }
275- }
276- if (actualExpertTotalNum < expertNum) {
277- expertCountOutLocal.SetValue((actualExpertTotalNum - 1) * NUM_TWO + NUM_TWO, 0);
278- expertCountOutLocal.SetValue((actualExpertTotalNum - 1) * NUM_TWO + NUM_THREE, 0);
279- }
280- 
281- event_t eventIDSToMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3));
282- SetFlag<HardEvent::S_MTE3>(eventIDSToMte3);
283- WaitFlag<HardEvent::S_MTE3>(eventIDSToMte3);
284- expertIdxCountOutQueue.EnQue<int64_t>(expertCountOutLocal);
285- sortedExpertIdxInQueue.FreeTensor(expertCountTempInLocal);
286-}
287- 
288-__aicore__ inline void FfnWbGroupListingArch35::expertCountCopyOut()
289-{
290- LocalTensor<int64_t> expertCountOutLocal = expertIdxCountOutQueue.DeQue<int64_t>();
291- DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(expertNum * NUM_TWO * sizeof(int64_t)),
292- 0, 0, 0};
293- DataCopyPad(groupListGm, expertCountOutLocal, copyParams);
294- 
295- expertIdxCountOutQueue.FreeTensor(expertCountOutLocal);
296-}
297- 
298-__aicore__ inline void FfnWbGroupListingArch35::CopyInOneCore(int64_t inputNum)
299-{
300- LocalTensor<int32_t> inLocal = sortedExpertIdxInQueue.AllocTensor<int32_t>();
301- DataCopyExtParams dataCopyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(inputNum * sizeof(int32_t)), 0, 0,
302- 0};
303- DataCopyPadExtParams<int32_t> dataCopyPadParams{false, 0, 0, 0};
304- DataCopyPad(inLocal, expandedExpertIdsGm, dataCopyParams, dataCopyPadParams);
305- sortedExpertIdxInQueue.EnQue<int32_t>(inLocal);
306-}
307- 
308-__aicore__ inline void FfnWbGroupListingArch35::ComputeOneCore(int64_t inputNum)
309-{
310- LocalTensor<int32_t> inLocal = sortedExpertIdxInQueue.DeQue<int32_t>();
311- SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
312- if (lastExpertId == -1) {
313- lastExpertId = inLocal.GetValue(0);
314- firstExpertId = lastExpertId;
315- }
316- for (int64_t i = 0; i < inputNum; i++) {
317- int32_t curExpertId = inLocal.GetValue(i);
318- tokenCount++;
319- if (curExpertId > lastExpertId && curExpertIdOffset < expertNum) {
320- // 先更新expertId切换前的数据
321- groupListOutLocal.SetValue(curExpertIdOffset * NUM_TWO, lastExpertId);
322- groupListOutLocal.SetValue(curExpertIdOffset * NUM_TWO + 1, tokenCount - 1);
323- curExpertIdOffset += 1;
324- tokenCount = 1;
325- lastExpertId = curExpertId;
326- }
327- }
328- sortedExpertIdxInQueue.FreeTensor(inLocal);
329-}
330- 
331-} // namespace FfnWbBatchingArch35
332-#endif // OP_KERNEL_FFN_WB_GROUP_LISTING_H
Dffn/ffn_worker_batching/op_kernel/arch35/ffn_wb_scan_sort_multi_core_arch35.h+0-517
@@ -1,517 +0,0 @@
1-/**
2- * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
3- *
4- * Licensed under the Apache License, Version 2.0 (the "License");
5- * you may not use this file except in compliance with the License.
6- * You may obtain a copy of the License at
7- *
8- * http://www.apache.org/licenses/LICENSE-2.0
9- *
10- * Unless required by applicable law or agreed to in writing, software
11- * distributed under the License is distributed on an "AS IS" BASIS,
12- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13- * See the License for the specific language governing permissions and
14- * limitations under the License.
15- */
16- 
17-/*!
18- * \file ffn_wb_scan_sort_multi_core_arch35.h
19- * \brief
20- */
21-#ifndef OP_KERNEL_FFN_WB_SCAN_SORT_MULTI_CORE_H
22-#define OP_KERNEL_FFN_WB_SCAN_SORT_MULTI_CORE_H
23-#include "ffn_wb_arch35_reuse.h"
24- 
25- 
26-namespace FfnWbBatchingArch35 {
27-using namespace AscendC;
28- 
29-class FfnWbScanSortMultiCoreArch35 : public SortMaskBase {
30-public:
31- __aicore__ inline FfnWbScanSortMultiCoreArch35(){};
32- __aicore__ inline void Init(GM_ADDR tokenInfoGm, GM_ADDR workspace, SortCustomTilingDataKernel *tilingData,
33- const ScheduleContextInfo *contextInfo, TPipe *tPipe);
34- __aicore__ inline void Process();
35- 
36-private:
37- __aicore__ inline void VBSProcess();
38- __aicore__ inline void UBSortProcess(int64_t progress, int64_t size, int64_t sortNum, int64_t loopSessionCnt);
39- __aicore__ inline void OneCoreVMSProcess(int64_t listNum, int64_t perListElements, int64_t lastListElements);
40- __aicore__ inline void VMSProcess();
41- __aicore__ inline void SortOutProcess();
42- __aicore__ inline void VBSCopyInAndClear(int64_t progress, int64_t size, int64_t sortNum, int64_t loopSessionCnt);
43- __aicore__ inline void UBSortCompute(int64_t progress, int64_t size, int64_t sortNum);
44- __aicore__ inline void VBSCopyOut(int64_t progress, int64_t size, int64_t sortNum);
45- __aicore__ inline void InitSortMaskMrgSort(SortCustomMrgsort *sorter, int64_t listNum, int64_t coreOffset,
46- int64_t sortNumCoreOffset, int64_t loopOffset, int64_t loopIdxOffset);
47- __aicore__ inline void InitSortMaskMrgSortOut(SortCustomMrgsortOut *sorter, int64_t listNum, int64_t coreOffset);
48- __aicore__ inline void CopyOutValidCount();
49- __aicore__ inline void ClearTokenInfoFlag();
50- 
51-private:
52- GlobalTensor<float> workspaceGms[NUM_TWO];
53- GlobalTensor<int32_t> workspaceSortNumGm_;
54- GlobalTensor<int32_t> expertIdsGmFStart_;
55- 
56- SortCustomTilingDataKernel *tilingData_ = nullptr;
57- const ScheduleContextInfo *contextInfo_ = nullptr;
58- 
59- int32_t totalValidCnt_ = 0;
60- int32_t curValidCnt_ = 0;
61- 
62- int64_t F = 0;
63- int64_t BsKLenWithPading = 0; // BS*K_plus_1按block对齐后的个数
64- 
65- // tilingData->perCoreSessionNum: 主核 每个核处理A中的几个
66- // tilingData->lastCoreSessionNum: 尾核 处理A中的几个
67- int64_t perCoreSortNum = 0; // perCoreSessionNum * BsKLenWithPading 主核 每个核排序的元素总个数; 不保证32个数对齐
68- int64_t lastCoreSortNum = 0; // lastCoreSessionNum * BsKLenWithPading;
69- int64_t sessionLoops = 0; // 当前核 ub循环次数
70- 
71- int64_t perLoopSessionNum = 0; // 当前核 一次ub处理A中的几个
72- int64_t lastLoopSessionNum = 0; // 当前核 尾ub处理A中的几个
73- int64_t perLoopElements = 0; // perLoopSessionNum * BsKLenWithPading;
74- int64_t lastLoopElement = 0; // lastLoopSessionNum * BsKLenWithPading;
75- 
76- // for MoeMrgsort
77- SortCustomMrgsort mrgsorter;
78- SortCustomMrgsortParam mrgsortParam;
79- 
80- int64_t blockIdx_ = 0;
81- int64_t srcWsIndex = 0;
82- int64_t bufferSize_ = 0;
83- 
84- int64_t listNum = 0;
85- int64_t perListElements = 0;
86- int64_t lastListElements = 0;
87- int64_t vmsSortNumStride_ = 0; // 核间
88- 
89- // int64_t sortTotalLength;
90- // int64_t sortCoreLoops;
91- // int64_t sortCoreLoopElements;
92- // int64_t sortCoreLastLoopElements;
93- 
94- static constexpr int64_t MAX_MRGSORT_LIST = 4;
95-};
96- 
97-__aicore__ inline void FfnWbScanSortMultiCoreArch35::VBSCopyInAndClear(int64_t progress, int64_t size, int64_t sortNum,
98- int64_t loopSessionCnt)
99-{
100- LocalTensor<int32_t> inLocal = sortDataCopyInQueue.AllocTensor<int32_t>();
101- // 根因修复(脏数据 bug)：AllocTensor 的 UB 未初始化，RECV 连跑复用残留脏值。整块预置 sentinel：
102- // expert 段(fp32 视图)MIN_FP32(<-expertStart_ 恒判无效)、rowIds 段 0；再由 DataCopyPad 覆盖真值。
103- // 保证尾块 padding 与 Cast/Compare 对齐 roundup 读区恒为确定值，杜绝脏尾污染 GatherMask 计数。
104- LocalTensor<float> inLocalFp32 = inLocal.ReinterpretCast<float>();
105- Duplicate<float>(inLocalFp32, MIN_FP32, sortNum);
106- Duplicate<int32_t>(inLocal[sortNum], 0, sortNum);
107- SetWaitFlag<HardEvent::V_MTE2>(HardEvent::V_MTE2);
108- int64_t inOffset = progress * perLoopSessionNum * contextInfo_->M * this->F;
109- DataCopyExtParams dataCopyParams{
110- static_cast<uint16_t>(loopSessionCnt),
111- static_cast<uint32_t>(contextInfo_->BS * contextInfo_->K * sizeof(int32_t)),
112- static_cast<uint32_t>((contextInfo_->M * this->F - contextInfo_->BS * contextInfo_->K) * sizeof(int32_t)), 0,
113- 0};
114- DataCopyPadExtParams<int32_t> dataCopyPadParams{true, 0, static_cast<uint8_t>(contextInfo_->BsKPaddingCount),
115- INT_MAX};
116- DataCopyPad(inLocal[0], expertIdsGm[inOffset], dataCopyParams, dataCopyPadParams);
117- 
118- LocalTensor<int32_t> rowIdsLocal = inLocal[sortNum];
119- // 索引生成只有(A, BsKPad)
120- int64_t startValue =
121- blockIdx_ * tilingData_->perCoreSessionNum * BsKLenWithPading + progress * perLoopSessionNum * BsKLenWithPading;
122- // SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
123- ArithProgression<int32_t>(rowIdsLocal, startValue, 1, size); // size包括: BsKLenWithPading
124- sortDataCopyInQueue.EnQue(inLocal);
125- 
126- // clear expertIds
127- LocalTensor<int32_t> clearLocal = sortDataCopyOutQueue.AllocTensor<int32_t>();
128- Duplicate<int32_t>(clearLocal, INT_MAX, size);
129- sortDataCopyOutQueue.EnQue(clearLocal);
130- 
131- clearLocal = sortDataCopyOutQueue.DeQue<int32_t>();
132- DataCopyExtParams copyoutParams{
133- static_cast<uint16_t>(loopSessionCnt),
134- static_cast<uint32_t>(contextInfo_->BS * contextInfo_->K * sizeof(int32_t)), 0,
135- static_cast<uint32_t>((contextInfo_->M * this->F - contextInfo_->BS * contextInfo_->K) * sizeof(int32_t)), 0};
136- SetWaitFlag<HardEvent::MTE2_MTE3>(HardEvent::MTE2_MTE3);
137- DataCopyPad(expertIdsGm[inOffset], clearLocal, copyoutParams);
138- sortDataCopyOutQueue.FreeTensor(clearLocal);
139-}
140- 
141-__aicore__ inline void FfnWbScanSortMultiCoreArch35::UBSortCompute(int64_t progress, int64_t size, int64_t sortNum)
142-{
143- LocalTensor<int32_t> inLocal = sortDataCopyInQueue.DeQue<int32_t>();
144- LocalTensor<int32_t> expertIdsLocal = inLocal[0];
145- LocalTensor<float> expertIdsLocalFp32;
146- 
147- expertIdsLocalFp32 = expertIdsLocal.ReinterpretCast<float>();
148- Cast(expertIdsLocalFp32, expertIdsLocal, RoundMode::CAST_ROUND, size);
149- 
150- // gathermask start
151- LocalTensor<uint32_t> maskLocalTensor = sortedBuffer.Get<uint32_t>();
152- uint64_t rsvdCnt = 0;
153- 
154- // 根因修复(脏数据 bug)：mask 缓冲(sortedBuffer)跨 loop/调用不清零，GatherMask 读到残留脏 mask 位
155- // 会把无效槽计入 → rsvdCnt 偏大/越界。用前整块清零，CompareScalar 只覆盖有效区、其余恒 0。
156- Duplicate<uint32_t>(maskLocalTensor, static_cast<uint32_t>(0),
157- Ceil(sortNum, ONE_REPEAT_COMPARE_NUM) * ONE_REPEAT_COMPARE_NUM / ONE_REPEAT_SORT_NUM);
158- 
159- Muls(expertIdsLocalFp32, expertIdsLocalFp32, (float)-1, size);
160- 
161- LocalTensor<uint8_t> maskLocalTensorUInt8 = maskLocalTensor.ReinterpretCast<uint8_t>();
162- AscendC::CompareScalar(maskLocalTensorUInt8, expertIdsLocalFp32, static_cast<float>(-expertStart_),
163- AscendC::CMPMODE::GT,
164- (size + ONE_REPEAT_COMPARE_NUM - 1) / ONE_REPEAT_COMPARE_NUM * ONE_REPEAT_COMPARE_NUM);
165- 
166- GatherMaskParams gatherMaskParams;
167- gatherMaskParams.repeatTimes = 1;
168- gatherMaskParams.src0BlockStride = 1;
169- gatherMaskParams.src0RepeatStride = 8; // 8 blocks
170- gatherMaskParams.src1RepeatStride = 0;
171- GatherMask(expertIdsLocalFp32, expertIdsLocalFp32, maskLocalTensor, true, size, gatherMaskParams, rsvdCnt);
172- curValidCnt_ = rsvdCnt;
173- if (rsvdCnt == 0) {
174- sortDataCopyInQueue.FreeTensor(inLocal);
175- return;
176- }
177- this->totalValidCnt_ += rsvdCnt;
178- int64_t duplicateNum = rsvdCnt % ONE_REPEAT_SORT_NUM;
179- if (duplicateNum > 0) {
180- int duplicateIndex = rsvdCnt - duplicateNum;
181- uint64_t mask0 = UINT64_MAX;
182- mask0 = mask0 << duplicateNum;
183- mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
184- uint64_t mask[NUM_TWO] = {mask0, 0};
185- Duplicate(expertIdsLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE);
186- }
187- int32_t selectedCnt = (rsvdCnt + ONE_REPEAT_SORT_NUM - 1) / ONE_REPEAT_SORT_NUM * ONE_REPEAT_SORT_NUM;
188- 
189- // step2: sort
190- LocalTensor<uint32_t> rowIdsLocal = inLocal[sortNum].ReinterpretCast<uint32_t>();
191- GatherMask(rowIdsLocal, rowIdsLocal, maskLocalTensor, true, size, gatherMaskParams, rsvdCnt);
192- LocalTensor<float> concatLocal = expertIdsLocalFp32;
193- LocalTensor<float> sortedLocal = sortedBuffer.Get<float>(GetSortLen<float>(sortNum));
194- LocalTensor<float> outLocal = sortDataCopyOutQueue.AllocTensor<float>();
195- Sort<float, true>(outLocal, concatLocal, rowIdsLocal, sortedLocal, selectedCnt / ONE_REPEAT_SORT_NUM);
196- 
197- sortDataCopyOutQueue.EnQue<float>(outLocal);
198- sortDataCopyInQueue.FreeTensor(inLocal);
199-}
200- 
201-__aicore__ inline void FfnWbScanSortMultiCoreArch35::VBSCopyOut(int64_t progress, int64_t size, int64_t sortNum)
202-{
203- int64_t oneLoopMaxSessionNum =
204- contextInfo_->sortLoopMaxElement / (contextInfo_->BS * contextInfo_->K + contextInfo_->BsKPaddingCount);
205- int64_t perCoreLoop = CeilDiv(tilingData_->perCoreSessionNum, oneLoopMaxSessionNum);
206- int64_t lastCoreLoop = CeilDiv(tilingData_->lastCoreSessionNum, oneLoopMaxSessionNum);
207- 
208- if (curValidCnt_ > 0) {
209- LocalTensor<float> outLocal = sortDataCopyOutQueue.DeQue<float>();
210- DataCopyExtParams copyParams{1, static_cast<uint32_t>(GetSortLen<float>(curValidCnt_) * sizeof(int32_t)), 0, 0,
211- 0};
212- int64_t wkOffset =
213- blockIdx_ * GetSortLen<float>(perCoreSortNum) + GetSortLen<float>(progress * perLoopElements);
214- 
215- if (perCoreLoop != lastCoreLoop && blockIdx_ == tilingData_->needCoreNum - 1 && lastCoreLoop == 1) {
216- DataCopyPad(workspaceGms[1][wkOffset], outLocal, copyParams);
217- } else {
218- DataCopyPad(workspaceGms[0][wkOffset], outLocal, copyParams);
219- }
220- 
221- sortDataCopyOutQueue.FreeTensor(outLocal);
222- }
223- 
224- // 个数为0，也要设置到gm上，防止脏数据
225- LocalTensor<int32_t> tempTensor = tempBuffer.Get<int32_t>(BLOCK_BYTES / sizeof(int32_t));
226- tempTensor.SetValue(0, curValidCnt_);
227- SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
228- DataCopyExtParams copyParams1{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(int32_t)), 0, 0, 0};
229- DataCopyPad(workspaceSortNumGm_[blockIdx_ * tilingData_->sortNumWorkSpacePerCore + progress], tempTensor,
230- copyParams1);
231-}
232- 
233-__aicore__ inline void FfnWbScanSortMultiCoreArch35::InitSortMaskMrgSort(SortCustomMrgsort *sorter, int64_t listNum,
234- int64_t coreOffset, int64_t sortNumCoreOffset,
235- int64_t loopOffset, int64_t loopIdxOffset)
236-{
237- GlobalTensor<float> srcWsGm = workspaceGms[srcWsIndex][blockIdx_ * coreOffset + loopOffset];
238- GlobalTensor<int32_t> srcSortNumGm = workspaceSortNumGm_[blockIdx_ * sortNumCoreOffset + loopIdxOffset];
239- LocalTensor<float> inLocal = sortDataCopyInQueue.AllocTensor<float>();
240- LocalTensor<float> outLocal = sortDataCopyOutQueue.AllocTensor<float>();
241- for (int64_t i = 0; i < listNum; i++) {
242- LocalTensor<float> inLocalT = inLocal[GetSortLen<float>(tilingData_->oneLoopMaxElementsMrg) * i];
243- sorter->SetInput(srcWsGm, srcSortNumGm, inLocalT);
244- }
245- GlobalTensor<float> dstWsGm = workspaceGms[1 - srcWsIndex][blockIdx_ * coreOffset + loopOffset];
246- LocalTensor<int32_t> outSortNumLocal = tempBuffer.Get<int32_t>(BLOCK_BYTES / sizeof(int32_t));
247- sorter->SetOutput(dstWsGm, outLocal, outSortNumLocal);
248- sortDataCopyInQueue.FreeTensor(inLocal);
249- sortDataCopyOutQueue.FreeTensor(outLocal);
250- tempBuffer.FreeTensor(outSortNumLocal);
251-}
252- 
253-__aicore__ inline void FfnWbScanSortMultiCoreArch35::InitSortMaskMrgSortOut(SortCustomMrgsortOut *sorter,
254- int64_t listNum, int64_t coreOffset)
255-{
256- GlobalTensor<float> srcWsGm = workspaceGms[srcWsIndex];
257- GlobalTensor<int32_t> srcSortNumGm = workspaceSortNumGm_;
258- LocalTensor<float> inLocal = sortDataCopyInQueue.AllocTensor<float>();
259- LocalTensor<float> outLocal = sortDataCopyOutQueue.AllocTensor<float>();
260- 
261- for (int64_t i = 0; i < listNum; i++) {
262- LocalTensor<float> inLocalT = inLocal[GetSortLen<float>(tilingData_->oneLoopMaxElementsMrg) * i];
263- sorter->SetInput(srcWsGm, srcSortNumGm, inLocalT);
264- }
265- 
266- LocalTensor<float> outLocalV = outLocal[tilingData_->oneLoopMaxElementsMrg * MAX_MRGSORT_LIST];
267- sorter->SetOutput(this->sortedexpertIdsGm, this->sortedRowIdsGm, outLocal, outLocalV);
268- 
269- LocalTensor<float> useTempBuffer =
270- sortedBuffer.Get<float>(GetSortLen<float>(tilingData_->oneLoopMaxElementsMrg) * MAX_MRGSORT_LIST);
271- sorter->SetBuffer(useTempBuffer);
272- sortDataCopyInQueue.FreeTensor(inLocal);
273- sortDataCopyOutQueue.FreeTensor(outLocal);
274-}
275- 
276-__aicore__ inline void FfnWbScanSortMultiCoreArch35::OneCoreVMSProcess(int64_t listNum, int64_t perListElements,
277- int64_t lastListElements)
278-{
279- int64_t coreOffset = GetSortLen<float>(perCoreSortNum);
280- int64_t sortNumCoreOffset = tilingData_->sortNumWorkSpacePerCore;
281- mrgsortParam.oneLoopMaxElements = tilingData_->oneLoopMaxElementsMrg;
282- 
283- int64_t curSortNumStride = 1;
284- for (int64_t i = 0; listNum >= 1; i++) {
285- int64_t loops = (listNum + MAX_MRGSORT_LIST - 1) / MAX_MRGSORT_LIST;
286- int64_t remainListNum = listNum - (loops - 1) * MAX_MRGSORT_LIST;
287- 
288- mrgsortParam.perListElements = perListElements;
289- mrgsortParam.sortNumStride = curSortNumStride;
290- 
291- int64_t loopOffset = GetSortLen<float>(mrgsortParam.perListElements * MAX_MRGSORT_LIST);
292- int64_t loopIdxOffset = mrgsortParam.sortNumStride * MAX_MRGSORT_LIST;
293- for (int64_t loop = 0; loop < loops - 1; loop++) {
294- InitSortMaskMrgSort(&mrgsorter, MAX_MRGSORT_LIST, coreOffset, sortNumCoreOffset, loop * loopOffset,
295- loop * loopIdxOffset);
296- mrgsorter.Init(&mrgsortParam);
297- mrgsorter.Process();
298- }
299- 
300- // now not use: mrgsortParam.perListElements = perListElements;
301- InitSortMaskMrgSort(&mrgsorter, remainListNum, coreOffset, sortNumCoreOffset, (loops - 1) * loopOffset,
302- (loops - 1) * loopIdxOffset);
303- mrgsorter.Init(&mrgsortParam);
304- mrgsorter.Process();
305- 
306- listNum = loops;
307- perListElements = perListElements * MAX_MRGSORT_LIST;
308- curSortNumStride = curSortNumStride * MAX_MRGSORT_LIST;
309- srcWsIndex = (srcWsIndex + 1) % WORK_GM_NUM;
310- if (loops == 1) {
311- break;
312- }
313- }
314-}
315- 
316-__aicore__ inline void FfnWbScanSortMultiCoreArch35::UBSortProcess(int64_t progress, int64_t size, int64_t sortNum,
317- int64_t loopSessionCnt)
318-{
319- VBSCopyInAndClear(progress, size, sortNum, loopSessionCnt);
320- UBSortCompute(progress, size, sortNum);
321- VBSCopyOut(progress, size, sortNum);
322-}
323- 
324-__aicore__ inline void FfnWbScanSortMultiCoreArch35::VBSProcess()
325-{
326- if (blockIdx_ < tilingData_->needCoreNum) {
327- int64_t sortNum = Ceil(perLoopElements, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
328- for (int64_t loop = 0; loop < sessionLoops - 1; loop++) {
329- UBSortProcess(loop, perLoopElements, sortNum, perLoopSessionNum);
330- }
331- 
332- sortNum = Ceil(lastLoopElement, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
333- UBSortProcess(sessionLoops - 1, lastLoopElement, sortNum, lastLoopSessionNum);
334- 
335- CopyOutValidCount();
336- 
337- if (sessionLoops > 1) {
338- OneCoreVMSProcess(sessionLoops, perLoopElements, lastLoopElement);
339- }
340- }
341- SyncAll();
342-}
343- 
344-__aicore__ inline void FfnWbScanSortMultiCoreArch35::VMSProcess()
345-{
346- int64_t currentStageNeedCoreNum = tilingData_->needCoreNumMrg;
347- perListElements = perCoreSortNum;
348- // now not use: lastListElements = lastCoreSortNum;
349- listNum = tilingData_->needCoreNum;
350- vmsSortNumStride_ = tilingData_->sortNumWorkSpacePerCore;
351- 
352- for (; listNum > MAX_MRGSORT_LIST;) {
353- currentStageNeedCoreNum = Ceil(listNum, MAX_MRGSORT_LIST);
354- int64_t coreOffset = GetSortLen<float>(perListElements * MAX_MRGSORT_LIST);
355- int64_t sortNumCoreOffset = vmsSortNumStride_ * MAX_MRGSORT_LIST;
356- int64_t remainListNum = listNum - (currentStageNeedCoreNum - 1) * MAX_MRGSORT_LIST;
357- 
358- mrgsortParam.perListElements = perListElements;
359- mrgsortParam.sortNumStride = vmsSortNumStride_;
360- mrgsortParam.oneLoopMaxElements = tilingData_->oneLoopMaxElementsMrg;
361- 
362- if (blockIdx_ < currentStageNeedCoreNum - 1) {
363- InitSortMaskMrgSort(&mrgsorter, MAX_MRGSORT_LIST, coreOffset, sortNumCoreOffset, 0, 0);
364- mrgsorter.Init(&mrgsortParam);
365- mrgsorter.Process();
366- } else if (blockIdx_ == currentStageNeedCoreNum - 1) {
367- InitSortMaskMrgSort(&mrgsorter, remainListNum, coreOffset, sortNumCoreOffset, 0, 0);
368- mrgsorter.Init(&mrgsortParam);
369- mrgsorter.Process();
370- }
371- listNum = currentStageNeedCoreNum;
372- // now not use: currentStageNeedCoreNum = Ceil(listNum, MAX_MRGSORT_LIST);
373- srcWsIndex = (srcWsIndex + 1) % WORK_GM_NUM;
374- 
375- // now not use: lastListElements = perListElements * (remainListNum - 1) + lastListElements;
376- perListElements = perListElements * MAX_MRGSORT_LIST;
377- vmsSortNumStride_ = vmsSortNumStride_ * MAX_MRGSORT_LIST;
378- 
379- SyncAll();
380- }
381-}
382- 
383-__aicore__ inline void FfnWbScanSortMultiCoreArch35::SortOutProcess()
384-{
385- if (blockIdx_ < 1) {
386- mrgsortParam.perListElements = perListElements;
387- mrgsortParam.sortNumStride = vmsSortNumStride_;
388- // now not use: mrgsortParam.lastListElements = lastListElements;
389- mrgsortParam.oneLoopMaxElements = tilingData_->oneLoopMaxElementsMrg;
390- 
391- SortCustomMrgsortOut sorter;
392- InitSortMaskMrgSortOut(&sorter, listNum, GetSortLen<float>(perListElements));
393- sorter.Init(&mrgsortParam, pipe);
394- sorter.Process();
395- }
396- SyncAll();
397-}
398- 
399-__aicore__ inline void FfnWbScanSortMultiCoreArch35::CopyOutValidCount()
400-{
401- LocalTensor<int32_t> outLocal = sortDataCopyOutQueue.AllocTensor<int32_t>();
402- // 根因修复(脏数据 bug)：outLocal 的 UB 未初始化。整块(1*32B=8*int32)清零、仅 [0] 写计数，
403- // 再以 32B 对齐 atomic-add：workspace[0]+=计数、[1..7]+=0(Init 已置 0)保持不变。
404- // V(memset)->S(SetValue)->MTE3(搬出) 逐级同步。
405- constexpr int32_t CNT_BLK_ELEM = static_cast<int32_t>(BLOCK_BYTES / sizeof(int32_t));
406- Duplicate<int32_t>(outLocal, static_cast<int32_t>(0), CNT_BLK_ELEM);
407- SetWaitFlag<HardEvent::V_S>(HardEvent::V_S);
408- outLocal.SetValue(0, totalValidCnt_);
409- SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
410- SetAtomicAdd<int32_t>();
411- DataCopy(rsvdCntGm[0], outLocal, CNT_BLK_ELEM);
412- SetAtomicNone();
413- sortDataCopyOutQueue.FreeTensor(outLocal);
414-}
415- 
416-__aicore__ inline void FfnWbScanSortMultiCoreArch35::Init(GM_ADDR tokenInfoGm, GM_ADDR workspace,
417- SortCustomTilingDataKernel *tilingData,
418- const ScheduleContextInfo *contextInfo, TPipe *tPipe)
419-{
420- this->pipe = tPipe;
421- tilingData_ = tilingData;
422- contextInfo_ = contextInfo;
423- 
424- F = contextInfo_->BS * contextInfo_->K + 1 + 1;
425- this->totalLength = tilingData_->totalLengthWithPad; // 这里最好改成 totalLength, 后面写workspace的地方也需要修改
426- 
427- BsKLenWithPading = contextInfo_->BS * contextInfo_->K + contextInfo_->BsKPaddingCount;
428- 
429- blockIdx_ = GetBlockIdx();
430- 
431- int64_t curCoreSessionNum = 0;
432- if (blockIdx_ == tilingData_->needCoreNum - 1) {
433- curCoreSessionNum = tilingData_->lastCoreSessionNum;
434- } else {
435- curCoreSessionNum = tilingData_->perCoreSessionNum;
436- }
437- 
438- perCoreSortNum = tilingData_->perCoreSessionNum * BsKLenWithPading;
439- lastCoreSortNum = tilingData_->lastCoreSessionNum * BsKLenWithPading;
440- 
441- int64_t oneLoopMaxSessionNum = contextInfo_->sortLoopMaxElement / BsKLenWithPading; // 大于0
442- 
443- sessionLoops = CeilDiv(curCoreSessionNum, oneLoopMaxSessionNum);
444- perLoopSessionNum = Min(oneLoopMaxSessionNum, curCoreSessionNum);
445- lastLoopSessionNum = curCoreSessionNum - (sessionLoops - 1) * perLoopSessionNum;
446- 
447- perLoopElements = perLoopSessionNum * BsKLenWithPading;
448- lastLoopElement = lastLoopSessionNum * BsKLenWithPading;
449- 
450- int64_t expertIdStartPos = contextInfo_->curMicroBatchID * F + 1 + 1;
451- expertIdsGmFStart_.SetGlobalBuffer((__gm__ int32_t *)tokenInfoGm + contextInfo_->curMicroBatchID * F);
452- expertIdsGm.SetGlobalBuffer((__gm__ int32_t *)tokenInfoGm + expertIdStartPos +
453- blockIdx_ * tilingData_->perCoreSessionNum * contextInfo_->M * F);
454- // rsvdCntGm 在scan阶段已经清零
455- rsvdCntGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace), SCAN_BATCHID_GM_OFFSET);
456- workspaceSortNumGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace) + OFFSET_SORTED_EXPERT_IDS,
457- contextInfo_->sortNumWorkSpace);
458- 
459- sortedexpertIdsGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace) + OFFSET_SORTED_EXPERT_IDS +
460- contextInfo_->sortNumWorkSpace,
461- this->totalLength);
462- sortedRowIdsGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace) + OFFSET_SORTED_EXPERT_IDS +
463- contextInfo_->sortNumWorkSpace + this->totalLength,
464- this->totalLength);
465- 
466- workspaceGms[0].SetGlobalBuffer((__gm__ float *)workspace + OFFSET_SORTED_EXPERT_IDS +
467- contextInfo_->sortNumWorkSpace + this->totalLength * NUM_TWO,
468- this->totalLength * NUM_TWO);
469- workspaceGms[1].SetGlobalBuffer((__gm__ float *)workspace + OFFSET_SORTED_EXPERT_IDS +
470- contextInfo_->sortNumWorkSpace + this->totalLength * (NUM_TWO + NUM_TWO),
471- this->totalLength * NUM_TWO);
472- 
473- bufferSize_ = Ceil(Max(tilingData_->oneLoopMaxElementsMrg * MAX_MRGSORT_LIST, contextInfo_->sortLoopMaxElement),
474- ONE_REPEAT_SORT_NUM) *
475- ONE_REPEAT_SORT_NUM * sizeof(int32_t) * NUM_TWO;
476- pipe->InitBuffer(sortDataCopyInQueue, bufferNum, bufferSize_);
477- pipe->InitBuffer(sortDataCopyOutQueue, bufferNum, bufferSize_);
478- pipe->InitBuffer(sortedBuffer, bufferSize_);
479- pipe->InitBuffer(tempBuffer, bufferSize_);
480-}
481- 
482-__aicore__ inline void FfnWbScanSortMultiCoreArch35::ClearTokenInfoFlag()
483-{
484- // 用最后一个核清理flag. 一个block一个有效数(0, int32_t). 总共需要A个block.
485- if (blockIdx_ == contextInfo_->coreNum - 1) {
486- int64_t perLoopElement = bufferSize_ / BLOCK_SIZE; // buffer总共可以支持的block个数
487- int64_t loops = Ceil(contextInfo_->A, perLoopElement);
488- int64_t lastLoopElement = contextInfo_->A - (loops - 1) * perLoopElement;
489- int64_t duplicateNum = Min(static_cast<int64_t>(contextInfo_->A), perLoopElement) * 8; // 8: block num
490- 
491- LocalTensor<int32_t> clearLocal = tempBuffer.Get<int32_t>();
492- Duplicate<int32_t>(clearLocal, 0, duplicateNum);
493- 
494- SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
495- 
496- int64_t curElementA = perLoopElement;
497- for (int64_t idx = 0; idx < loops; idx++) {
498- if (idx == loops - 1) {
499- curElementA = lastLoopElement;
500- }
501- 
502- DataCopyExtParams copyOutParams{static_cast<uint16_t>(curElementA), static_cast<uint32_t>(sizeof(int32_t)),
503- 0, static_cast<uint32_t>((contextInfo_->M * F - 1) * sizeof(int32_t)), 0};
504- DataCopyPad(expertIdsGmFStart_[idx * perLoopElement * contextInfo_->M * F], clearLocal, copyOutParams);
505- }
506- }
507-}
508- 
509-__aicore__ inline void FfnWbScanSortMultiCoreArch35::Process()
510-{
511- VBSProcess();
512- ClearTokenInfoFlag();
513- VMSProcess();
514- SortOutProcess();
515-}
516-} // namespace FfnWbBatchingArch35
517-#endif // OP_KERNEL_FFN_WB_SCAN_SORT_MULTI_CORE_H