PyPTO-Gym 算子设计 R0 阶段：Module 划分方法论与实战指南

PyPTO-Gym 算子设计 R0 阶段：Module 划分方法论与实战指南

【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym

Module 划分是 PyPTO-Pro 算子 tile 级方案设计（pypto-pro-op-design）中 R0 轮的核心任务：在动手规划 Tile、循环与片上布局之前，先把整体计算流拆成逻辑清晰、可独立验证的计算模块。本文以 module_partitioning.md 为骨架，结合 PyPTO-Gym 仓库中的设计流程、脚本与模板，系统讲解如何根据数据依赖与Section 边界完成 Module 划分，并产出可供下游 R1–R8 轮直接消费的 Module 契约。读完本文，你将掌握稳定 Softmax、Online Softmax、Cube+Vector 组合算子等典型场景的 Module 切分方法，并能正确填写 R0 的 Module 列表与module_interfaces.yaml。

Module 是什么：逻辑划分，不是语法结构

Module 是设计文档中对计算流程的逻辑划分，不是 PyPTO Pro 的语法结构。换句话说，Module 只存在于 DESIGN.md 与module_interfaces.yaml这类设计产物中，kernel 源码里并不存在名为 "Module" 的代码块。这一点决定了划分时的两个基本视角：

• 数据视角：计算流中哪些步骤之间存在"必须先完成 A 才能开始 B"的依赖关系；
• 资源视角：相邻步骤合并在同一个 Module 后，UB 等片上空间能否放下同时存活的数据。

Module 划分主要看数据依赖和Section 边界，同时需要评估片上资源是否足够。整个 R0 轮在 pypto-pro-op-design 的九轮迭代设计（R0–R8）中承担"分解计算流"的职责，其完整流程见 SKILL.md 的「R0：Module 划分」一节；而 R0 的产出将写入 design-template.md 的 §0「Module 划分（R0 输出）」以及机器可读的module_interfaces.yaml。

Module 划分按两步进行：

1. 根据数据依赖找出必须分开的计算步骤；
2. 根据API 和数据通路判断 Section 边界：Cube 和 Vector 计算属于不同 Section，也要拆成不同 Module。

下面分别展开这两步。

第一步：根据数据依赖确定边界

对于相邻的两步计算 A 和 B，是否拆成两个 Module，可以按以下常见情况判断：

情形	判断	理由
B 必须等 A 遍历完完整的归约轴（或其他完整数据范围），并基于 A 的最终结果开始一次新的遍历或独立计算阶段	通常划分为两个 Module	B 的输入是 A 的全局结果，二者无法在同一轮循环内完成
A 处理完当前 Tile 后，B 能在同一层循环内立即消费该 Tile，并且两步属于同一个 Section	通常放在同一个 Module	数据可在循环内直接接力，无需跨遍历保存状态
A 和 B 之间只有少量同域计算，没有插入其他 Section 或独立计算阶段	可以放在同一个 Module	依赖紧、边界不明显，拆分开反而增加设计复杂度
B 只是紧接着对 A 得到的归约结果做汇总或简单处理，不需要重新遍历完整数据范围	可以继续放在 A 所在的 Module	B 不依赖对完整数据范围的重遍历

需要特别强调的是，这些是常用判断方式，不是硬性规则。有两个容易误解的点：

• A 的结果是否需要保存，不能单独决定 Module 边界。同一 Module 内也可能暂存中间结果（例如跨 Tile 循环的累加状态），不同 Module 之间也可能通过片上空间直接传递数据，并不一定落 GM。
• 即使 A 和 B 可以连续执行，如果合并后 UB 等片上空间放不下同时存活的数据，也需要拆分 Module。资源约束可以反向制造边界——这是"也需要注意片上资源是否足够"的落点。

案例一：沿 N 轴分块的稳定 Softmax（三遍遍历）

以沿 N 轴分块的稳定 Softmax 为例。常规实现需要三次完整遍历，后两次遍历都依赖前一次得到的最终结果，因此通常按三次遍历划分三个 Module：

Module 1：遍历全部 N Tile，得到每行的 m = max(x) Module 2：重新遍历全部 N Tile，得到每行的 s = sum(exp(x - m)) Module 3：重新遍历全部 N Tile，写出 y = exp(x - m) / s

对同一个行块，前一个 Module 遍历完全部 N Tile 后，下一个 Module 才能开始。这里的边界来自"基于最终结果重新遍历 N 轴"这一数据依赖，不是因为m或s需要保存——m、s作为跨 Module 的归约状态，即使在同一个 Module 内部也需要持久化，可见"是否需要保存"与"是否拆分 Module"是两个维度。

该三遍结构与 api_mapping_and_tile_planning.md 中给出的稳定 Softmax 数据依赖图一致：x → reduce_max(N) → max，随后subtract → exp → reduce_sum(N) → sum，最后exp ÷ divide → out。其中max、sum的 shape 均为[M, 1]，后续计算需要沿 N 轴广播。

案例二：Online Softmax（两遍遍历）

采用 Online Softmax 时，可以在一次遍历中同时更新最大值和指数和。处理完前 k 个 N Tile 后，每行维护两个跨 Tile 状态：

• m：当前最大值；
• s：以当前m为基准的指数和，即sum(exp(x - m))。

每次合入一个新 Tile 的更新公式为：

m_new = max(m, tile_max) s_new = s * exp(m - m_new) + sum(exp(tile - m_new))

由于最终的m和s到遍历结束时才能确定，写输出时通常要重新读取各 Tile，因此典型的 Module 划分是两遍：

Module 1：在线遍历全部 N Tile，得到最终 m 和 s Module 2：重新遍历全部 N Tile，写出 exp(x - m) / s

这里有一个重要的边界细节：如果后续只对最终的m或s做一次汇总或简单运算，不再重新遍历 N 轴，这部分可以留在 Module 1 中。这与"根据数据依赖确定边界"表格中第四条判断（B 只是紧接着对 A 的归约结果做汇总，可留在同一 Module）完全对应。

关于 Online Softmax 的完整递推与尾块处理，仓库中的模式文档 online-softmax-tail.md 给出了更完整的公式（含输出累加量o）与数值安全要求：状态m、l、o的累加应保持在 FP32，尾块 mask 必须在row_max之前生效（"padding 零对 max 不是恒等元"），并且 PyPTO-Pro 中模块级float("inf")无法编译（会渲染成未声明的 C++ 标识符inff），应使用有限哨兵值NEG_LARGE = -1.0e30。

案例三：Attention 的在线归一化——不能直接套用 Softmax 公式

Attention 的在线归一化与 Online Softmax 的关键区别在于：除了m和s，它还要维护输出累加量。m更新时，已有的输出累加量也要按新的m缩放，再合入当前 Tile。因此：

Attention 应按完整的在线更新公式划分 Module，不能直接套用上面的 Softmax 公式。

完整的 Attention 在线更新递推（见 online-softmax-tail.md 的 Complete recurrence）为：

m_chunk = row_max(s) m_new = max(m_old, m_chunk) alpha = exp((m_old - m_new) * scale) p = exp((s - m_new) * scale) l_new = l_old * alpha + row_sum(p) o_new = o_old * alpha + p @ v_chunk

最终y = o_new / l_new。其中分母更新不可分割：只乘alpha不加row_sum(p)是错误的。这条模式文档同时提示：QK/PV 交接与 cross-core 同步属于目标版本相关的内容，应从官方代码推导，而不是从递推公式照搬——这与 Module 划分后由 R6 轮（跨核同步）承接的工作边界一致。

第二步：根据 Section 边界划分

数据依赖分析完成后，再根据API 使用的执行域、内存空间和硬件通路判断相邻步骤之间是否存在 Section 边界。Section 是 PyPTO Pro 中按执行域组织的代码结构（pl.section_cube()/pl.section_vector()），不同 Section 使用不同的硬件流水，因此跨 Section 的数据交接天然形成 Module 边界。

各数据通路对应的 Section 归属如下：

计算/搬运路径	所属 Section	说明
L1/L0A/L0B/L0C 上的矩阵路径	Cube Section	包括 GM→L1、L1→L0A/L0B、矩阵乘和 L0C 写回
Vec（UB）上的数据搬运和向量计算	Vector Section	使用 VF 时，外层 Kernel 负责 GM 与 UB 之间的搬运，@pl.vector_function负责 UB 与寄存器之间的数据交换和寄存器计算
pl.quant和pl.dequant	Vector Section	走 V 流水，输入、scale、offset 和输出均为 Vec Tile
pl.move/pl.store对 Acc 结果做随路量化或反量化	Cube Section	操作走 FIX 流水；per-channel 参数使用MemorySpace.ScalingTile 时，参数先从 Mat（L1）搬到 Scaling，再由 FIX 使用
UB→L1、L0C→UB 等跨执行域搬运	按当前move接口文档选择数据通路	以目标版本 API 文档为准

关于量化参数缓冲有一个易踩的坑：Scaling是量化参数使用的片上缓冲区，不是独立的 Section。不能看到MemorySpace.Scaling就新建 Module；仍要根据使用它的 API 和数据通路判断属于 Cube 还是 Vector。同理，pl.quant/pl.dequant因为走 V 流水放在 Vector Section，而pl.move/pl.store对 Acc 的随路量化走 FIX 流水放在 Cube Section——同样是量化，归属完全不同。

一个 Module 只能属于一个 Section。相邻步骤分别属于 Cube 和 Vector 时，在两者之间划分 Module。R0 只记录每个 Module 属于 Cube 还是 Vector；具体使用几个 Section 代码块、多个同域 Module 是否放在同一个 Section，以及循环相对 Section 的位置，在 R4（循环与 Section 结构）确定。这符合 loop_design.md 的约定：R0 已经记录每个 Module 属于 Cube 还是 Vector，R4 据此确定源码中的 Section 结构，相邻的同域 Module 可以顺序写在同一个 Section 中。

案例四：Cube 与 Vector 组合算子

Matmul 后接 Vector 后处理是融合算子的典型形态，Section 边界同时形成 Module 边界：

Cube Module：A × B → intermediate ↓ Vector Module：intermediate → 激活、归约或归一化 → output

两个 Module 可以由同一次 Kernel 启动完成。中间数据可以写到 GM workspace，也可以走move接口支持的片上通路。R0 记录数据的写入位置、读取位置和同步方向；具体同步点和event_id在 R6 确定，规则见 cross_core_synchronization.md。

这类"整数 Cube 收缩 + FP 浮点 Vector 尾处理"的融合形态，在仓库的 cv-quant-matmul-direct-epilogue.md 模式文档中有详细约束：整数收缩在 K 全过程中保持 int32，再在同一 kernel 内进入 FP 浮点 Vector epilogue；epilogue 顺序（整数 pre-bias → FP32 scale → 可选 offset → 浮点 post-bias → 最终 cast）必须从 golden 冻结。由于 Cube 与 Vector 分属不同 Module，跨执行域的数据交接（Acc 直搬 Vec，或经 GM workspace）与 READY/FREE 信用同步必须按跨核同步规则设计。

关于 Cube 路径的数据通路，api_mapping_and_tile_planning.md 给出了完整的链路：pl.matmul(dst, lhs, rhs)要求lhs位于 Left（L0A）、rhs位于 Right（L0B）、dst位于 Acc（L0C），一次矩阵乘通常展开为GM → Mat(L1) → Left/Right(L0A/L0B) → matmul/matmul_acc → Acc(L0C) → store/move → GM 或 Vec(UB)。VF 路径则呈现GM → Vec(UB) → vf.load* → 寄存器 → vf.* → vf.store* → Vec(UB) → GM的层次。这些通路信息正是第二步判断 Section 归属时的硬件依据。

R0 输出：Module 列表与 module_interfaces.yaml

R0 的产出落在两个地方：DESIGN.md 的 §0，以及机器可读的module_interfaces.yaml契约。

Module 列表

每个 Module 记录数学目标、Section、输入、输出和前置依赖：

Module	数学目标	Section	输入	输出	前置依赖
Module 1	...	Cube/Vector	...	...	...

只有一个 Module时，必须说明数据依赖和 Section 边界为何不需要继续拆分（例如纯逐元素、单遍完成且不跨执行域的算子）。跨 Tile 状态及其循环范围在 R4 补充——即 loop_design.md 中"结果单元"与"跨 Tile 状态生命周期"的职责。

DESIGN.md 模板 §0 还要求记录归约轴容量结论：归约轴是否可能超单 tile、多 Tile 归约方案（常规多遍 / 在线统计加输出两遍 / 其他），以及 Module 级数据流示意。这些信息直接呼应 R0 的划分依据。

module_interfaces.yaml 契约

R0 阶段必须产出机器可读的module_interfaces.yaml（single source of truth），包含以下字段：

• module_count：Module 总数；
• is_fusion：是否为融合算子（同时含 cube 和 vec section →true）。按 R0 定义，is_fusion=true隐含module_count >= 2（Cube 和 Vector 必须划分到不同 Module）；
• has_cross_core：是否涉及 cross_core 跨核流水（信息记录用，不影响分流判据）；
• modules[]：每个 Module 的id/name/description/section（cube或vector）/golden_steps/inputs（source 为primary或module_<j>，且j < 当前 id）/outputs/golden_stage_fn；
• final_outputs：每个 golden 返回值对应到产出 Module；
• composition_verification：atol / rtol / seeds / shapes。

骨架可由脚本自动生成：

python ./scripts/gen_module_interfaces.py custom/<op>/<op>_golden_cpu.py \ --spec custom/<op>/SPEC.md --op <op> \ --design custom/<op>/DESIGN.md > custom/<op>/module_interfaces.yaml

脚本 gen_module_interfaces.py 会从{op}_golden_cpu.py中解析 golden 函数签名，自动填充schema_version、op、primary_inputs与composition_verification等确定性部分；modules[]边界、final_outputs接线、is_fusion、has_cross_core等标注TODO的判断部分由 architect 手工填写。

产出后必须自验：

python ./scripts/validate_module_yaml.py custom/<op>/module_interfaces.yaml --json

返回"status": "PASS"才算完成。验证器 validate_module_yaml.py 内置了完整的契约规则（rule0–rule7）：文件必须是声明了op/module_count/modules的非空映射且至少一个 Module（防止生成器错误消息被重定向进 YAML 后空文件通过校验）；inputs[*].source == "primary"的名称必须存在于primary_inputs；source == "module_j"必须满足j < 当前 id且名称存在于module_j.outputs（禁止前向/自引用）；final_outputs的module_j必须满足j <= module_count且名称存在于对应 Module 输出；不允许 no-op Module；shape 表达式只允许+ - * //与名称/整数 token；dtype 必须在允许的词汇表内；module_count必须等于len(modules)。该脚本还支持--self-test运行全部规则的自检用例，其中记录了从旧拼写phase_<j>到规范拼写module_<j>的演进，两套拼写目前均被接受以保证历史产物继续通过校验。

Module 划分与后续轮次的衔接

Module 划分不是孤立的设计动作，它与后续轮次存在明确的上下游关系：

• R1（API 映射）：在 R0 的 Module 划分基础上，将 API 调用链细化到 Module 内部每一步操作；
• R4（循环与 Section）：R4不重新划分 Module，只负责把 R0 的 Module 落实到 Section 与循环结构；若 API 数据通路与 R0 记录的 Section 归属冲突，需回到 R0 修正 Module 划分；
• R6（跨核同步）：Cube 与 Vector 之间的数据交接使用set_cross_core/wait_cross_core，R0 记录数据的写入位置、读取位置和同步方向，R6 落实具体同步点与event_id；
• R8（综合评估）："数据依赖是否正确（Module 顺序、sync 位置）"与"跨 tile 状态是否正确初始化和持久化"均可能回溯到 R0 重新评估。

其中最关键的一条回溯链是：归约轴超单 tile 时回 R0。当 R7.5 目标测试 case 验证发现"归约轴超单 Tile 的跨 tile 状态与持久化"适配不了时，要回溯 R0 重新设计归约方案（常规多遍 / 在线统计加输出两遍），而不是删改测试 case 迁就设计。这再次印证了 R0 Module 划分质量对整条设计链的决定性影响。

小结

Module 划分的两把尺子是：数据依赖决定"必须分开"，Section 边界决定"按执行域分开"，同时片上空间容量可以强制制造额外的拆分边界。判断边界时始终记住：是否需要保存中间结果不构成划分依据；Scaling不是 Section；一个 Module 只能属于一个 Section。对 Softmax 类多遍遍历算子，边界来自"基于最终结果重新遍历归约轴"；对 Cube+Vector 融合算子，Section 边界就是 Module 边界，跨域交接的同步细节交给 R6。最终把划分结论同时落进 DESIGN.md §0 的 Module 列表与机器可读的module_interfaces.yaml，并通过 validate_module_yaml.py 的 PASS 校验，R0 才算真正闭环。

【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym