pypto-gym 中的 linspace 组合实现arange mul 逐 tile 生成等差序列的 NPU kernel 骨架【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym本文以 linspace.md 这篇 kernel 参考骨架为主体讲清torch.linspace在 PyPTO 框架下的组合实现方案为什么 linspace 不是原子接口、arange/full/mul/assemble四个 API 如何拼出等差序列、无 batch 轴时沿输出轴 tile loop 切分这一模式的设计动机以及骨架中各占位符的含义与适用边界。读完后你能够照骨架落地一个生成类算子的 kernel 草稿并知道它距离可上 NPU 还需要经过哪些验证环节。1. 背景linspace 在 Torch ↔ Pypto 映射表中的定位pypto-gym 的 cannbot-skills/ops/pypto-api-explore/references/torch-pypto-op-mapping.md 维护了一张 Torch 与 Pypto 的对标表。其中「同名映射」部分明确列出Pypto 的创建类原子接口只有arange、full、ones、zeros见该文件第 19 行创建分组linspace并不在原子接口之列因此被归入「组合方案」Torch 算子Pypto 组合方案参考实现linspacearangemullinspace.md对照同组的eyearangeone_hotcast见 eye.md可以看出 PyPTO 处理生成类算子的思路是一致的不直接造数而是用 arange 生成索引序列再经逐元素运算与写入操作组合出目标数据。linspace 的目标序列out[i] start i * (end - start) / (n - 1)中start为 0 的典型情形如模型里生成0..num_grid_per_side-1的网格坐标就退化为i * step恰好只需要arangemul两步这就是组合方案表里只写两个 API 的原因。2. 参考骨架完整代码与逐行解读原骨架给出的 kernel 如下来自 linspace.md文档中的 Note 一句话点明了切分方式无 batch 轴沿输出轴按 tile loop 切分每 tile 生成一段等差序列。pypto.frontend.jit(runtime_options{run_mode: pypto.RunMode.NPU}) def linspace_kernel(out: pypto.Tensor(ol, pypto_dtype)): for i in pypto.loop(num_tiles, nameout_tile, unroll_list[1]): step pypto.full([tile_len], 1.0 / max(n - 1, 1), pypto_dtype) idx pypto.cast(pypto.arange(tile_len), pypto_dtype) seg pypto.mul(idx, step) pypto.assemble(seg, [i * tile_len], out)2.1 接口签名只有一个输出张量函数签名out: pypto.Tensor(ol, pypto_dtype)说明本算子没有输入张量ol输出 shape 列表与pypto_dtype元素 dtype都是外部预先算好的量。这一结论与 examples 目录的占位符约定一致examples/README.md 定义了全部骨架共用的占位符与本骨架直接相关的是占位符含义ol输出 shape 列表pypto_dtype元素 dtype如pypto.DT_FP32n输出总长度linspace 的元素个数num_tiles/tile_len生成类沿输出轴切 tile 的分块数 / 每块长度linspace 等生成类算子的所有输入其实都是标量参数起止值、元素个数 nshape 本身由 host 侧确定这正是骨架里没有输入pypto.Tensor参数的原因。2.2 tile loop输出轴切分的动机与写法pypto.loop(num_tiles, nameout_tile, unroll_list[1])沿输出轴把整个等差序列切成num_tiles段、每段tile_len个元素即n num_tiles * tile_len。切分的收益在 README 的约定里有明确说法骨架要展示哪些轴 loop、哪些轴整块的接口组合模式。对生成类算子而言逐 tile 生成可以控制单次 kernel 迭代的工作量与 UB统一缓冲占用避免一次性把整个大张量放进片上处理同时 loop 内所有计算只依赖i当前 tile 序号和常量host 侧即可完全确定不存在数据依赖的动态轴问题。unroll_list[1]表示该循环以 1 为步长、不做循环展开与同目录其它骨架如 cat.md 的unroll_list[1]保持一致具体tile_len取多少、要不要展开README 已声明需按实际 shape / dtype 与平台约束确定并调优。2.3 tile 内三步计算full → arangecast → mul每轮迭代生成第i段序列写法上有三个值得注意的细节step pypto.full([tile_len], 1.0 / max(n - 1, 1), pypto_dtype)步长1/(n-1)是 host 侧标量这里用pypto.full把它物化成长度为tile_len的常量张量。max(n - 1, 1)是显式的除零防护n 1时退化输出只有一个元素分母取 1 而非 0与torch.linspace在numel1时不再除步长的行为对齐。idx pypto.cast(pypto.arange(tile_len), pypto_dtype)arange生成的是整数索引序列0..tile_len-1创建类原子接口见 torch-pypto-op-mapping.md 第 19 行随后cast到目标 dtype 再参与乘法避免整数乘浮点步长时的类型问题。seg pypto.mul(idx, step)等差序列即索引 × 步长。骨架对应的是 start0 的通用形态若 start 非 0按组合方案再补一次add逐元素add属同名同参原子接口即可。这里体现的是 PyPTO 生成类 kernel 的通用数据流arange 负责位置full 负责常量逐元素运算负责取值assemble 负责落位。2.4 assemble 落位元数据式写入pypto.assemble(seg, [i * tile_len], out)把当前 tile 写到输出张量偏移i * tile_len处。assemble在 PyPTO 中是view 取片 assemble 拼回的搬运语义——同目录 cat.md 的 Note 对这类算子有同样的定性本算子为 metadata/搬运语义view 取片 assemble 拼回无逐元素计算。linspace 的 loop 内只有构造与写入没有跨张量归约因此多 tile 之间互不干扰不需要 pypto-specific-ops.md 中atomic_add一类的多核合并手段。从源码结构看仓库内生产实现也大量使用assemble按 tile 偏移写回输出例如 deepseek_v2_lite_chat/mla_prolog/mla_prolog.py 中按[tIdx, 0, 0]偏移分段 assemble以及 deepseek_v32_exp/lightning_indexer_quant_impl.py 中按s1_tile_idx * s1_tile偏移写回与本骨架i * tile_len的偏移模式是同一套路。3. 与同类骨架的横向对照何时用 batch loop何时用输出轴 tile loopexamples 目录下每个op.md都遵循 README 的 Note 约定用一句话说明切分方式且无 batch-row loop 的骨架cube、sort、生成/索引类须说明实际采用的切分方式。linspace 正属于生成类与 eye.md 的对照能说明选轴逻辑eye有隐含的行轴输出第 0 轴沿行 loop每行整块生成 one-hotlinspace输出只有一维没有 batch 轴可借只能沿唯一的输出轴自己切 tilecat有显式 batch 轴sl[0]沿 batch loop 搬运。即 loop 轴的选取原则是有 batch 轴优先 loop batch纯生成/索引类则沿输出轴切 tile。这一点也解释了骨架中num_tiles、tile_len这对占位符只出现在生成类沿输出轴切 tile的场景见占位符表第 18 行而不是所有骨架通用。4. 生产模型里的真实需求linspace 用在哪骨架对应的现实需求可以在仓库的 transformers 代码中直接找到。以 qwen3_5_9b/modeling_qwen3_5.py 为例视觉塔插值模块按grid_thw生成归一化网格坐标for _, h, w in grid_thw_list: h_idxs torch.linspace(0, self.num_grid_per_side - 1, h) w_idxs torch.linspace(0, self.num_grid_per_side - 1, w) h_idxs_floor h_idxs.int() w_idxs_floor w_idxs.int()qwen3_6_27b 的 modeling 文件 中有完全相同的用法。这类调用的特点是start 恒为 0终点是网格边长减一元素个数 h/w 由输入图片决定——正是骨架里 start0 形态、n由 host 侧标量给出的形态序列长度不大但属于无输入张量、纯参数驱动的生成类算子切分维度天然只有输出轴。从源码结构看这两处目前仍直接调用torch.linspacePyPTO 侧要承接该语义组合方案arange full mul tile loop assemble就是现成的实现路径。5. 使用骨架的正确姿势探索流程、门禁与验证边界5.1 在 API 探索工作流中的位置linspace.md 属于 pypto-api-explore 技能包的一部分。按该 SKILL 的工作流命中「组合方案」条目后应按映射表取用组合方案linspace →arangemul并阅读 examples 下对应骨架占位符约定以 examples/README.md 为准README 给出了最小可运行 setup 示例B, D 8, 128等命中条目后仍需通过 Explore 核实具体约束与生产实现再按 templates/api_report.md 生成含「API 映射 / 约束 / Tiling / 风险评估 / 证据索引」等章节的报告门禁要求 6 个章节非空。5.2 骨架的边界声明必须遵守README 对 examples 目录有两条硬声明引用骨架时不能越过骨架仅展示接口组合与轴切分模式……不作为标准模板loop 轴、unroll_list、tile shape、动态轴处理等需按实际 shape / dtype 与平台约束确定并调优所有骨架未逐一经 NPU 编译验证。因此正确流程是以骨架为草稿 → 按实际n/dtype 设定num_tiles、tile_len并确认 tile 数与 tile 长度符合 set_vec_tile_shapes 相关 Tiling 约束 中每维 0、最多 4 维等硬约束→ 走仓库内 pypto-op-verify、pypto-kernel-validate 等技能包定义的测试流程做精度比对后再使用。此外注意run_mode: pypto.RunMode.NPU表明该骨架面向 NPU 运行模式骨架中1.0 / max(n - 1, 1)的 host 侧除法要求n在编译期已知n为动态值时需要在风险评估中按 SKILL 文档动态 shape 兼容性一节标注并调整切分策略。6. 小结要点结论依据linspace 无原子接口归入「组合方案」arangemultorch-pypto-op-mapping.md切分模式无 batch 轴沿输出轴 tile loopnum_tiles × tile_len nlinspace.md除零防护max(n - 1, 1)处理n 1退化情形骨架源码落位方式assemble(seg, [i * tile_len], out)搬运语义、tile 间无数据竞争骨架源码及 cat.md 同类 Note生产对应场景ViT 网格坐标生成start0 形态modeling_qwen3_5.py使用边界骨架非标准模板、未逐一经 NPU 编译验证需按实际约束调优并走验证流程examples/README.md一句话概括本文主线linspace 在 PyPTO 中组合实现 arange 造索引、full 造步长、mul 出序列、tile loop 控规模、assemble 按i * tile_len落位骨架给出的是接口组合与切分模式的最小表达实际落地前需按占位符约定填参、按 SKILL 门禁核实约束、按验证流程完成精度比对。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考