PyPTO 逐模块实现指南从 DESIGN 到可验证 Kernel 的完整开发流程【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym本文基于 module-development.md 展开面向在 CANN / pypto-gym 仓库中使用 PyPTO 框架开发自定义算子的开发者与 Agent。文章系统讲解单模块与多模块两种实现路径、实现前的核对清单、验证与失败定位方法、pypto.frontend.jit的 Kernel 配置、循环与 Tile 设计、数据读取写回方式、动态长度与状态管理以及数据类型与精度策略。读完本文你将掌握从读取 DESIGN.md 到产出通过detailed_tensor_compare精度验证的op_impl.py的完整工作流并能避开本仓库实测中反复出现的反模式。一、模块划分单模块与多模块两条路径PyPTO 算子的实现以 DESIGN.md 中的模块划分为基准。模块只用于组织开发与验证并不要求每个模块都有独立的 JIT kernel——所有模块最终会整合进同一个生产级pypto.frontend.jit入口。单模块路径L0当设计判定module_count 1时DESIGN.md §0.3复杂度预算total_complexity 1.3直接产出单个op_impl.py用完整 golden 做一次端到端验证即可。此路径不涉及分阶段文件链。编排规则 中明确L0 路径下 staged chain is skipped entirely — singleop_impl.pyis the deliverable。多模块路径L1module_count 2时每次只实现一个语义模块验证通过后再实现下一个。编排器使用累计实现文件推进op_module1_impl.py # 仅 M1 的真实 PyPTO 逻辑 op_module12_impl.py # M1 M2 op_module123_impl.py # M1 M2 M3 ... op_module1…N_impl.py # 全模块 端到端 kernel命名与调度约定遵循编排规则后缀数字是累计的1、12、123…每个文件都要通过detailed_tensor_compare对所有输出的比较才允许进入下一个阶段文件。每个文件的命名约定为op_modulek_wrapperlint 规则 OL08 强制。多模块路径的三条架构禁令见编排规则Prohibition A禁止一次性把所有模块拼成一个 JIT仅 L0 允许 one-shotProhibition B禁止用 Python hostfor循环模拟 kernel 的分块迭代逻辑Prohibition C禁止每模块一个生产 JIT的架构——模块是语义块不是独立 JIT 入口。验证中间模块时可以导出中间结果或直接使用 golden 提供的边界张量但被验证的计算必须与最终实现一致——不能为了通过边界检查而临时改变算法。二、实现准备动笔前必须核对的清单开始写代码前逐项核对以下内容源自module-development.md核对项说明API 签名与 DESIGN.md /module_interfaces.yaml中的模块契约一致输入/输出 shape、dtype与 golden 契约一致动态轴显式标注切片秩与偏移pypto.view的 offsets / valid_shape 落在原 Tensor 范围内尾块有效范围动态长度下最后一块可能不满 tile必须用valid_shape标记循环状态与 tile 配置状态变量作用域、TileShape 是否在对应操作前设置要点采用 DESIGN.md 中确定的Tile、显式配置和类型转换位置。设计文档中的候选值仅供后续比较不代表已验证可用关键值未确定时应先补全设计而不是猜。调整参数时遵守记录的限制见执行约束并重新验证受影响的结果。使用实现模板作为强制骨架模板覆盖 Layer G–KG cache/bridge、H pypto_* 子内核、I kernel 实现、J JIT 入口、K host wrapper并内嵌了 OL45 / OL46 / OL47 反模式注释结构规范详见 pypto-kernel-design-format.md 的 Layer A–L 定义。遇到具体问题再读调试手册而不是盲目试错。Golden 覆盖检查实现必须覆盖 golden 中属于本模块的每一个计算包括掩码、缩放、cast、状态更新和全部输出。编排规则第 19 条要求逐项对照 golden 头部注释中的函数清单Golden function inventory在 MEMORY.md 中给每个数学操作标记 ✅对应 pypto 调用 行号或 ❌缺失存在 ❌ 时不允许继续推进模块。三、验证与失败定位以证据驱动不做口头断言验证入口将实现文件路径和模块输出交给验证 Agent先检查编译与 shape/dtype再用detailed_tensor_compare与 golden 做数值比较。编译通过或结构检查通过都不能代替精度比较。detailed_tensor_compare位于 detailed_tensor_compare.py核心行为支持tuple / list / dict嵌套输出的递归比较对每个叶子张量都生成报告tensor_leaf_pairs会逐层展开shape 不匹配直接判all_closeFalse输出total_elements、out_of_tolerance_count、out_of_tolerance_ratio、max_diff、mean_diff、std_diff以及超出容差元素的前 20 个明细索引、两侧值、绝对/相对差容差默认rtol1e-3, atol1e-3可通过options或旧式位置参数调整对inf处理有专门语义两侧同为同符号inf视为匹配否则计入超差。调用方式由验证 runner 统一引入无需手动配置 PYTHONPATHfrom detailed_tensor_compare import detailed_tensor_compare result detailed_tensor_compare(kernel_out, golden_out, tensor_nameoutput) assert result[all_close], Mismatch in output编排规则强制每个阶段文件、每个输出tuple/list/dict 的每个叶子都必须比较禁止只挑一个输出验证。设计检查 vs 实现测试DESIGN.md 中的设计检查结论只说明已核对的文档、推导或接口不能直接当作实现测试结果。必须按其验证方案执行并记录实际证据编排规则第 4 条禁止用口头 should pass 代替真实运行。失败定位策略验证通过后记录该模块输出的比较结果继续下一个模块后续失败时先找出最早输出错误的模块检查其输入输出及中间值必要时二分定位不要因为后续模块失败就随意修改已验证的模块确需修改时重新验证该模块及所有依赖它的结果。二分定位可参考上板二分定位法在 kernel 函数中把检查点 tensor 作为输入参数支持原地修改与 golden 对应位置的中间结果比较从单个关键计算点开始逐步二分直到定位首个出错的 op。检查点必须与 golden 的 shape、dtype、数量、顺序完全一致循环内的变量可在循环外创建大 tensor循环内用view/assemble赋值。已知问题排查路径通用调试、精度调试、中间值比较。每个调整必须对应具体原因不能盲目叠加精度变通参数。收尾全部模块完成后验证组合及所有最终输出并确认临时输入或检查点没有进入生产计算路径cleanup 调度会整理出独立的op_impl.py。四、Kernel 配置与计算函数JIT 入口与 helper 拆分JIT 入口统一使用pypto.frontend.jit。计算较复杂或需要复用时把 PyPTO 计算放入普通 helper由 JIT 入口调用def compute(x, output): # 根据设计完成切分、加载、计算和写回。 ... pypto.frontend.jit def kernel(x, output): compute(x, output)注意helper 仍在构图环境中执行不能把它当作独立的 CPU 数值函数。另外存在一个已知陷阱——若 helper 体内使用pypto.is_loop_begin(idx)/pypto.is_loop_end(idx)parser 会在编译期抛F00002, ValueError: Not concrete value且无源行信息。两种规避方式见 impl_template.py.tmpl 与 SKILL.md 注意点 19推荐把整个 body 直接 inline 进pypto.frontend.jit入口替代给该 helper 加pypto.frontend.function仅支持 tensor 参数非 tensor 参数当前不支持。pass_options 与 runtime_options根据具体实现选择不要照搬固定取值参数用途cube_l1_reuse_setting配置 Cube L1 复用vec_nbuffer_setting、cube_nbuffer_setting配置缓冲副本数需计入资源预算涉及 UB 容量估算stitch_function_max_num配置子图合并数量device_sched_mode选择目标设备支持的调度方式生产级用法见 SKILL.md还包括ready_on_host_tensors: [cu_seqlens_q, ...]runtime_options标量索引 tensor 驻留 host避免 device→host 同步读取pypto.experimental.set_operation_options(combine_axisTrue)轴合并优化上述 setting 的 per-op keyed 形式如{0: 8, 1: 1}/{-2: 1, 0: 8, 1: 2}按算子索引精调优于全局{-1: x}sg_set_scopeN包裹 reduce / softmax 链融合为一次向量超算子中间量流式通过、不逐个物化 UB。语义标签需要识别不同计算段时可用pypto.set_semantic_label(scores)等有含义的名称。标签只用于定位代码段不替代正确的数据依赖和 Tile 配置——不要让标签看起来对就忽略了真正的依赖关系。五、循环与 Tile 设计循环选择与返回值遵循循环约束需要表达方式返回值按运行时边界重复执行pypto.loop符号索引按指定因子展开计算pypto.loop_unroll(索引, 展开因子)少量、编译期确定的迭代Pythonrange具体整数索引关键约束C-LOOP-* 规则节选C-LOOP-02符号循环边界必须用pypto.loop不能交给 PythonrangePython 无法把运行时符号转成整数C-LOOP-03符号循环索引不能用作 Python 容器下标或 Python 布尔条件C-LOOP-04循环内读取循环前尚未提交的计算结果时核对提交依赖——submit_before_loopTrue用于进入循环前提交已有任务不代表每轮迭代之间同步C-LOOP-05展开因子从单一候选开始loop_unroll解包索引与展开因子两个返回值C-LOOP-06跨迭代保留的状态在合适的外层作用域建立明确初始化、更新及最终写回C-LOOP-08loop_unroll覆盖被展开轴的数据访问必须以解包的展开因子确定尺寸如pypto.view(x, [uf, ...], [i, ...])不得写死常量尺寸。需要按展开因子调整每次处理的数据量时使用loop_unroll返回的因子。注意多个展开候选会增加编译路径——实现阶段unroll_list只能含单一值默认[1]照搬 DESIGN.md 「范式与设计决策」 中的选定值多值展开如{2, 1}仅允许在性能调优阶段使用lint OL56 强制 FAIL。Tile 约束详见Tiling 约束C-TILE-02尾轴 32B 对齐——tile_last × dtype_bytes % 32 0FP32 为 8 个元素FP16/BF16 为 16 个元素C-TILE-05矩阵乘 m/k/n 各轴使用[L0, L1]配置满足0 L0 L1且L1 % L0 0Vector 配置不能替代 Cube 配置C-TILE-06TileShape 必须在对应操作前设置参数必须是编译期整数或可解析为整数的常量不能来自运行时 shape、kernel 参数或 SymbolicScalarlint OL48 强制C-TILE-08计算形状变化时重新核对当前 Vector TileShape。Tile 作用域遵循 per-stage 原则pypto-kernel-design-format.md §11c多个pypto_*子内核需要不同 tile 时把set_cube_tile_shapes/set_vec_tile_shapes放到各自子内核内部而不是在_kernel_impl顶部设一个全局值OL47 会提示。六、数据读取与写回场景表达方式注意事项连续 tilepypto.view(tensor, shape, offsets, valid_shape...)显式指定有效范围时使用 view简单连续切片tensor[start:stop, :]核对目标版本的索引及有效形状支持分页 KVview、索引及结果拼装参考分页加载模式稀疏索引读取pypto.index_select核对索引轴、类型DT_INT32/DT_INT64和边界按模式选择元素pypto.gathermask例如 RoPE 奇偶位置核对 mode 的含义输出 tilepypto.assemble(tile, offsets, output)或切片赋值核对秩、偏移和输出有效区域Cache 更新pypto.scatter_update明确索引和重复写入的行为几个源码级要点pypto.view参数类型execution-constraints.md §5.2shape必须全是 Python int不接受 SymbolicScalaroffsets接受 SymbolicScalarvalid_shape接受 SymbolicScalar用于尾块有效数据标记。Python[]切片内部会对 index 做int()转换因此也不能用 SymbolicScalar 做切片索引。分页 KV 拼装AT-17在 loop 外分配拼装缓冲区loop 内按block_table逐块 view 拷贝零搬运尾块用valid_shape标记实际长度。禁止替换为gather_in_l1/gather_in_ub——后者是显式搬运指令机制冲突且引入真实开销。assemble写回pypto.assemble(tile, offsets, output)没有返回值直接修改outoffsets必须小于out.shape。同一 Tensor 在同一图里既被view读取、又被assemble写回会形成图成环报错见执行约束 §4.9。输出写回三选一OL02out[:] ...、out.move(...)、pypto.assemble(..., out)out ...只会绑定局部变量不会修改出参。scatter_update不支持 broadcastdim保持默认-2。七、动态长度与状态位置动态长度计算length tensor.shape[axis] tile_count (length tile_size - 1) // tile_size valid_len (length - offset).min(tile_size) # 需要将复杂符号表达式标记为中间变量时单独调用 seq_len cu_seqlens[batch 1] - cu_seqlens[batch] seq_len.as_variable()as_variable()原地修改符号对象、返回 None。它不负责分配 Tensor也不自动建立循环状态——只是把运行时标量显式标记为变量便于后续引用。注意valid_len中的.min()在 kernel 内对 SymbolicScalar 使用s.min(other)对 Tensor 逐元素使用pypto.minimumPython 原生min/max只能用于宿主侧逻辑执行约束 §4.5。状态变量的作用域状态典型位置原因Online Softmax 累积器Q tile 循环内、KV tile 循环外每个 Q tile 独立累积递推状态Batch 循环内、序列循环外每个序列保持独立状态单轮临时张量当前循环体内无跨迭代依赖需要全零状态时使用目标 API 支持的分配和填充方式如pypto.full注意fill_value与dtype必须一致动态图尾块无法自动推导有效范围时必须显式传valid_shape。必须在正确作用域初始化避免每轮重置。循环内累加的标准模式执行约束 §5.4acc pypto.tensor([TILE, D], pypto.DT_FP32, acc) for idx in pypto.loop(n, nameLOOP, idx_nameidx, unroll_list[1]): # 实现阶段单一值 tile compute_something(...) if pypto.is_loop_begin(idx): acc[:] tile # 首次迭代初始化 else: acc[:] acc tile # 后续迭代累加 if pypto.is_loop_end(idx): result pypto.cast(acc, pypto.DT_BF16) pypto.assemble(result, [offset, 0], output) # 最后一次写回注意pypto.tensor()创建的是未初始化随机值必须在is_loop_begin分支中初始化实现阶段高频陷阱 #1用pypto.full在循环前物化初始化会触发 F00003 对齐错误正确做法是pypto.tensor纯声明 is_loop_begin内 shape-matched 赋值。多动态轴2D reshape 嵌套 loop concrete tile当算子有 2 个及以上动态轴如 Batch SeqLen时不能直接在高维 tensor 上调受限 APImatmul 编译期需要 concrete shapeDYN 维度在编译期表现为 -1必须采用4D [B, N, S, D] ↓ 在 Python wrapper 层做 reshape 2D [B*N*S, D] ↓ 进入 kernel ↓ pypto.loop(b) → pypto.loop(N) → pypto.loop(s_tiles) ↓ pypto.view([S_TILE, D], [symbolic_offset, 0], valid_shape[actual_s, D]) 2D tile [S_TILE, D] ← shape 全是 concrete int ↓ matmul / elementwise / sum ↓ pypto.assemble(result, [symbolic_offset, 0], output_2d)关键点shape 全 concrete、动态性只进入 offset 和 loop bound、valid_shape 处理尾块、2D matmul 编译期完全确定。此外梯度算子多输出在不同维度累加时如 dQ 沿 S2 累加、dK/dV 沿 S1 累加使用两趟分离计算避免跨 loop 的读写依赖无需submit_before_loopTrue代价是中间结果重复计算一次。八、数据类型与精度在具体计算中标明输入、累加、转换和输出四种类型常见的浮点归约采用FP32 累加后转回输出类型API 约束 C-API-05精度敏感的归约和跨循环累加优先使用 FP32转换位置与参考计算的数值要求一致INT8 矩阵乘可产生 INT32 累加结果再反量化FP8 的缩放及转换位置取决于所选 API不能套用统一路径不同版本的 FP8 matmul API 对 scale 的处理方式不同必须查对应 API 文档。除法可按目标 API 选择pypto.PrecisionType.HIGH_PRECISION或INTRINSIC并与 golden 比较误差后决定。精度规则见 API 约束量化计算见对应的 AT 卡片patterns/atoms。其他 dtype 相关约束执行约束matmul显式给out_dtype调用前必须设置set_cube_tile_shapes(...)3D/4D 场景还要设置set_vec_tile_shapes(...)cast显式暴露CastMode和SaturationMode浮点转整数时satmodeON/OFF会直接改变溢出后的结果值sum只支持DT_FP32需要提高累加精度时显式转换为 FP32标量参与计算且 dtype 不能依赖隐式映射时使用pypto.Element(dtype, value)构造顺序固定。九、总结逐模块开发的纪律逐模块开发的核心纪律可以浓缩为四点一次一个模块多模块路径下任一时刻只有一个语义模块的 PyPTO 逻辑处于未冻结状态后续模块用 stub 或 golden 边界张量占位并在代码中显式注释# STUB: until Mk verified; golden-fed tensor证据驱动推进每个模块的边界验证都必须有detailed_tensor_compare的实际运行结果并记录到 MEMORY.md 的 Per-module verification loglint 失败不得判完成失败先定位再修改从最早出错的模块开始二分定位不因后续失败随意回改已验证模块收尾清理全部模块通过后确认临时输入/检查点没有进入生产计算路径再整理出独立的op_impl.py与README.md。遵循以上流程配合实现模板、执行约束与调试手册即可在 CANN / pypto-gym 仓库中稳定地完成从设计到可验证 Kernel 的算子开发闭环。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考