ARTICLE DETAIL

资讯详情

深耕商务建站与企业官网运营的一线实战洞察。

PyPTO-Gym 外积 kernel 参考:unsqueeze + mul 组合实现与按行 loop 切分实践

PyPTO-Gym 外积 kernel 参考:unsqueeze + mul 组合实现与按行 loop 切分实践 PyPTO-Gym 外积 kernel 参考unsqueeze mul 组合实现与按行 loop 切分实践【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym本文围绕 PyPTO-Gym 仓库中pypto-api-explore技能包下的 outer kernel 参考骨架讲解如何用 PyPTO 的unsqueezemul组合方案实现torch.outer外积算子并深入剖析其“沿第一输入元素输出行loop 切分、第二输入轴整块广播”的 Tiling 策略。读完本篇你可以理解 PyPTO 组合算子的 kernel 骨架写法、占位符约定与 view→compute→assemble 的切分范式并能将同一模式迁移到 RoPE 频率矩阵构建等真实模型算子场景。1. 算子映射outer 在 PyPTO 中的组合方案PyPTO 没有提供与torch.outer同名的原子接口。在仓库的算子对标手册 torch-pypto-op-mapping.md 中outer被明确归入「数学运算」类的组合方案Torch 算子Pypto 组合方案参考实现outerunsqueezemulouter.md从映射逻辑看外积out[i, j] a[i] * b[j]本质上是一次逐元素广播乘法把两个一维向量各自扩展出一个新轴a变成列向量、b变成行向量再用mul做广播相乘即可得到二维结果。这也决定了它属于纯 Vector 类算子仅逐元素运算不涉及 matmul实现上无需 Cube 侧的set_cube_tile_shapes重点在于 loop 切分与广播轴的处理。该骨架的取用方式由技能工作流定义pypto-api-explore/SKILL.md 的「步骤 2.5: 本地映射优先」要求——算子命中「命名映射」或「组合方案」条目时按表取用对应的 Pypto API 组合方案并阅读 examples/ 下对应的 kernel 参考骨架。2. outer kernel 完整参考骨架outer.md 给出的完整骨架如下占位符约定见下一节Note: 沿第一输入元素输出行loop 切分每行 a[i] × b 向量第二输入轴整块。pypto.frontend.jit(runtime_options{run_mode: pypto.RunMode.NPU}) def outer_kernel(a: pypto.Tensor(al, pypto_dtype), b: pypto.Tensor(bl, pypto_dtype), out: pypto.Tensor(ol, pypto_dtype)): b_row pypto.unsqueeze(b, 0) for i in pypto.loop(batch, namerow, unroll_list[1]): a_i pypto.view(a, [1, 1], [i, 0]) r pypto.mul(a_i, b_row) pypto.assemble(r, [i, 0], out)逐行拆解pypto.frontend.jit(runtime_options{run_mode: pypto.RunMode.NPU})以 PyPTO 前端 JIT 编译入口声明 kernelruntime_options指定 NPU 运行模式pypto.Tensor(al, pypto_dtype)kernel 参数按(shape 列表, dtype)声明al/bl/ol分别是输入a、输入b与输出out的 shape 占位符pypto_dtype为元素 dtype如pypto.DT_FP32b_row pypto.unsqueeze(b, 0)在 loop 之外对b提一次级补出新首轴使其成为行向量供循环内逐行广播复用pypto.loop(batch, namerow, unroll_list[1])沿输出行轴长度由占位符batch给出建立切分 loop每次迭代处理一行unroll_list[1]表示按 1 份展开a_i pypto.view(a, [1, 1], [i, 0])用view从a中按偏移[i, 0]取出一个1×1的 tile即标量a[i]r pypto.mul(a_i, b_row)1×1tile 与整块b行向量广播相乘得到输出的第i行pypto.assemble(r, [i, 0], out)把计算得到的行 tile 按偏移[i, 0]组装回输出张量。这套「view取 tile → 逐元素计算 →assemble写回」的三段式是 PyPTO loop 切分算子的通用数据流。仓库中的生产实现也采用同一模式例如 mla_prolog.py 中沿 tiling 轴view出子块、处理后assemble回输出sum_lstm.py 同样以view取 tile、assemble写回。3. 切分策略解析为什么按“行” loop、第二输入轴整块骨架开头的 Note 一句话点明了 Tiling 决策值得展开loop 轴 输出行即第一输入a的轴外积输出是二维的[len(a), len(b)]行与行之间相互独立第i行只依赖a[i]与整个b因此沿行轴 loop 切分是天然的并行切分方向每轮迭代的计算量为len(b)个元素的逐元素乘法整块轴 第二输入b的轴b对每一行都是完整参与的且是纯广播不需要逐元素索引变化无需再切整块参与可避免额外的切片开销a每次只取一个元素view(a, [1, 1], [i, 0])把a[i]变成1×1tile 后广播等价于把标量乘到整行。从源码结构看这种“一个 loop 轴 若干整块广播轴”的骨架形式与同目录其他参考如 rope.md 的 batch 轴 loop、cos/sin 同行整块参与乘加保持一致每篇骨架只用一句话说明“哪些轴 loop、哪些轴整块、为什么”。4. 占位符约定与最小可运行 setupexamples/README.md 定义了所有骨架共用的占位符体系阅读outer.md时须对照理解占位符含义sl输入 shape 列表如[B, S, D]ol输出 shape 列表pypto_dtype元素 dtype如pypto.DT_FP32batch被 loop 的外层轴长度通常sl[0]inner单次迭代处理的内层 shape如sl[1:]inner_out单次迭代的输出内层 shape末轴长度可能变化outer 骨架在此基础上额外使用了al/bl/ol三个 shape 列表占位符分别对应输入a、输入b与输出。结合骨架中view(a, [1, 1], [i, 0])、assemble(r, [i, 0], out)的偏移写法可以推断其心智模型为batch al[0]loop 的行数输出ol为二维外积形状b的第二轴整块参与。README 同时给出各骨架的最小可运行 setup 示例examples/README.md#L20-L29import pypto B, D 8, 128 sl, ol [B, D], [B, 1] pypto_dtype pypto.DT_FP32 batch, inner, inner_out B, [D], [1]将B、D具体化为 outer 场景时B对应a的长度loop 轴长度D对应b的长度整块广播轴长度。5. outer 的真实应用场景RoPE 频率矩阵外积不是抽象练习——仓库中多个模型的前向路径都在使用torch.outer构建 RoPE旋转位置编码的频率矩阵。以 Qwen3.5 视觉端 RoPE 实现 为例def forward(self, seqlen: int) - torch.Tensor: seq torch.arange(seqlen, deviceself.inv_freq.device, dtypeself.inv_freq.dtype) freqs torch.outer(seq, self.inv_freq) return freqs这里seq是长度为seqlen的位置索引向量inv_freq是长度为dim/2的逆频率向量外积得到[seqlen, dim/2]的频率矩阵。同一模式也出现在 OpenPangu Dense 模型、GutenOCR 的 Qwen2.5-VL 建模文件 等多处且测试代码 test_pangu_fused_layer.py 中同样以torch.outer(t, inv_freq)作为参考行为。这个用法与 outer 骨架的切分策略严丝合缝seqlen轴即 loop 的“输出行”轴每一行等于seq[i]标量×inv_freq整块向量。若在 NPU 侧将该步骤下沉为 PyPTO kernelouter kernel 参考 的行 loop 切分正是可直接对号入座的骨架。6. 使用边界与注意事项引用该骨架时务必注意 examples/README.md 的三条约定骨架非标准模板每个op.md仅为 kernel 参考骨架只展示接口组合与轴切分模式哪些轴 loop、哪些轴整块loop 轴、unroll_list、tile shape、动态轴处理等需按实际 shape / dtype 与平台约束确定并调优未经逐一经 NPU 编译验证所有骨架未逐一经 NPU 编译验证落地前需要走完整验证流程Note 约定每篇骨架以一句话说明切分方式loop 的轴、整块的轴及原因outer 的 Note 即“沿第一输入元素输出行loop 切分每行 a[i] × b 向量第二输入轴整块”。7. 延伸阅读完整的 API 探索工作流输入解析、公式分解、并行探索、报告生成见 pypto-api-explore/SKILL.mdTorch ↔ Pypto 全量对标同名映射、命名映射、组合方案三大类见 torch-pypto-op-mapping.md切分模式相近的参考骨架rope.mdbatch 轴 loop last-dim 折半旋转、norm.mdmulsumsqrt组合生产侧同范式的 view/assemble 切分实现mla_prolog.py、sum_lstm.py。综上outer在 PyPTO 中的落地路径是标准的“组合方案”unsqueeze构造广播轴 mul完成逐元素外积配合沿输出行轴的loop切分与view/assemble数据流即可覆盖 RoPE 频率矩阵构建等高频真实场景。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表
PREV
查看更多资讯
NEXT
返回资讯列表