已合并
feat(ffn_worker_batching/arch35): 下一代完整重写,解耦优化 #9960
feat(ffn_worker_batching/arch35): 下一代完整重写,解耦优化 #9960
已合并
zl_hw创建于 25 天前
zl_hw
25 天前

描述

FfnWorkerBatching算子纯下一代风格解耦重写

关联的Issue

https://gitcode.com/cann/ops-transformer/issues/4287

测试

ut/st

文档更新

不涉及

类型标签

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 zl_hw 的贡献)
Zzl_hw
25 天前 创建了 pull request,commit 19707918
atomgit-bot
atomgit-bot
25 天前 评论:

变更摘要

本次 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 结构体同步扩展,新增 flatElementssortSegNummergeRoundswsFlatIds 等字段描述各阶段切片与 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.hffn_wb_scan_sort_multi_core_arch35.hffn_wb_scan_sort_multi_core_bsk_arch35.hffn_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 中删除了 sortLoopMaxElementsortNumWorkSpace,新增 flatElementspreparePerLoopRowssortSegNumsortPerSegElementssortLenPerSegexpertStartmergeRoundsmergeOneLoopElementsextractPerLoopElementsglRowsPerLoop 以及 6 个 workspace 段偏移字段(wsFlatIdswsGatherIdx),kernrel 侧严格不自行计算任何布局。

  • Kernel 入口统一为模板函数ffn_worker_batching_apt.cpp 将原先分离的 FfnWorkerBatchingProcess(NORM)和 FfnWorkerBatchingRecvProcess(RECV)合并为单一模板 FfnWorkerBatchingA5<isRecv>,按五相位流水线顺序执行并在相位间插入 SyncAll;RECV 路径中增加了 ScheduleContextParse 后的 SyncAll 以防止轮询下标被提前推进。

likedislike
CANN-robotCANN-robot成员
25 天前 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
25 天前 评论:

CLA Signature Pass

zl_hw, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
CANN-robot
CANN-robot成员
25 天前 评论:

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 /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)
likedislike
此处折叠了56条消息 查看更多
唐超成员
24 天前 通过了评审
唐超成员
24 天前 评论:

/lgtm
/approve

likedislike
CANN-robotCANN-robot成员
24 天前 添加了label:lgtm
CANN-robotCANN-robot成员
24 天前 关闭了关联的issue
CANN-robotCANN-robot成员
24 天前 合入了pull request