Pull Request已成功合入, 合并人@CANN-robot
(感谢 zl_hw 的贡献)变更摘要
本次 PR 对 arch35(Ascend950 / DAV_3510)平台的 FfnWorkerBatching 算子进行了完整重写(A5),彻底解除了对 A2 实现的依赖。旧版通过 ffn_wb_arch35_reuse.h 桥接头复用 A2 的排序、归并、group_listing 等 11 个 kernel 头,新版本以五相位流水线(prepare → 段内排序 → 段间归并 → 拆包 → group_list → gather)完全自研实现。Host 侧 tiling 不再沿用 A2 的切分阈值与 workspace 公式,改为按运行时 UB 容量自行推导各阶段参数;kernel 侧不再自算任何布局或切分,所有参数由 host tiling 下发。TilingData 结构体同步扩展,新增 flatElements、sortSegNum、mergeRounds、wsFlatIds 等字段描述各阶段切片与 workspace 段偏移。
主要改动
-
Host Tiling 完全重写:
ffn_worker_batching_tiling_arch35.cpp中删除了 A2 路径的固定核数阈值(TH_RECV_CORE_NUM)、手工扣减UB_SYS_RESERVED_SIZE的做法,改为通过ReserveLocalMemory(RESERVED_SIZE_32K)声明预留后再取容量;核数由「工作量能否喂饱一个核」动态决定;workspace 布局由DoOpTiling统一逐段累加wsFlatIds_至wsGatherIdx_,GetWorkspaceSize直接取userWorkspaceWords_而非单独重算公式。 -
新增 A5 原生 kernel 头替代桥接复用:删除
ffn_wb_arch35_reuse.h及所有依赖 A2 的排序/group_listing 文件(ffn_wb_sort_multi_core_arch35.h、ffn_wb_scan_sort_multi_core_arch35.h、ffn_wb_scan_sort_multi_core_bsk_arch35.h、ffn_wb_group_listing_arch35.h),新增ffn_wb_a5_context.h(契约解析 + 工具函数,通过offsetof读取ScheduleContext)、ffn_wb_a5_prepare.h(RECV 忙等握手 + expert_id 归一化)、ffn_wb_a5_sort.h(段内 VBS 排序 → 段间 VMS 归并 → 拆包三合一)、ffn_wb_a5_group_list.h(基于比较掩码的向量游程编码,输出[expert_id, tokenNum]稠密行)、ffn_wb_a5_gather.h(按 gather_idx 取行搬运,含向量化整数除法索引换算)。 -
TilingData 结构体扩展:
ffn_worker_batching_arch35_tiling_def.h中删除了sortLoopMaxElement和sortNumWorkSpace,新增flatElements、preparePerLoopRows、sortSegNum、sortPerSegElements、sortLenPerSeg、expertStart、mergeRounds、mergeOneLoopElements、extractPerLoopElements、glRowsPerLoop以及 6 个 workspace 段偏移字段(wsFlatIds至wsGatherIdx),kernrel 侧严格不自行计算任何布局。 -
Kernel 入口统一为模板函数:
ffn_worker_batching_apt.cpp将原先分离的FfnWorkerBatchingProcess(NORM)和FfnWorkerBatchingRecvProcess(RECV)合并为单一模板FfnWorkerBatchingA5<isRecv>,按五相位流水线顺序执行并在相位间插入SyncAll;RECV 路径中增加了ScheduleContextParse后的SyncAll以防止轮询下标被提前推进。


Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here
PR Approval Progress
✅ Congratulations! All modules have met the lgtm and approve requirements.
Module Approval Details
| module | lgtm status | approve status |
|---|---|---|
| ffn | ✅ 马兵, 唐超 (2/2) | ✅ 唐超, 马兵 (2/1) |
💡 Tip:
- Committer can comment
/approveor/lgtm- Commenting
/approveimplies both code review (lgtm) and intent to merge (approve)


描述
FfnWorkerBatching算子纯下一代风格解耦重写
关联的Issue
https://gitcode.com/cann/ops-transformer/issues/4287
测试
ut/st
文档更新
不涉及
类型标签