PyPTO 中 split 算子的 view 切片实现：batch 轴 loop 搬运与轴切分骨架解析

PyPTO 中 split 算子的 view 切片实现batch 轴 loop 搬运与轴切分骨架解析【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym导读本文基于 CANN pypto-gym 仓库中 pypto-api-explore 技能的 split kernel 参考骨架系统讲解如何在 PyPTO 编程框架下实现split类轴切分算子batch 轴用pypto.loop循环搬运切分轴用pypto.view取一片view 切片其余轴整块搬运。读完本文你将掌握pypto.loop/pypto.view/pypto.set_vec_tile_shapes/pypto.assemble四个接口在轴切分场景下的组合方式理解占位符参数sl/ol/batch/inner/inner_out的语义并能结合仓库中的真实算子实现如 Arctic LSTM、Qwen3.5 GDR把骨架扩展为可运行的 NPU kernel。split 的 Torch→PyPTO 映射view 切片而非独立 APIPyPTO 框架没有为split提供独立的原子接口。仓库中的 Torch ↔ Pypto 算子对标手册 明确将split/chunk归类为「差异映射」条目split/chunk→view切片需按份数计算各片偏移也就是说torch.split(tensor, split_size_or_sections, dim)在 PyPTO 中需要由开发者手工完成两件事按份数计算各片的起始偏移offset这是split与普通view的关键差异用pypto.view从原张量上截取对应切片切片仍是原 GM 数据的视图不产生额外搬运。而由于 NPU 片上存储UB容量有限一个完整 shape 的切片往往无法一次放入因此 split 骨架采用了「batch 轴 loop、逐片搬运」的经典 Tiling 策略——这正是本文 split.md 骨架的核心思想。骨架代码全解四步完成一次切分split.md 给出的完整骨架如下注释为逐行解析pypto.frontend.jit(runtime_options{run_mode: pypto.RunMode.NPU}) def split_kernel(a: pypto.Tensor(sl, pypto_dtype), out: pypto.Tensor(ol, pypto_dtype)): # 第 1 步batch 轴 loop 搬运。i 为当前迭代的 batch 偏移 # unroll_list[1] 表示按 1 个 batch 为单位展开/流水 for i in pypto.loop(batch, namebatch, unroll_list[1]): # 第 2 步从输入 a 上取当前 batch 行 全部内层的切片 a_s。 # 偏移量 [i, 0, 0, ...]即 batch 轴偏移 i其余轴偏移 0整块 a_s pypto.view(a, [1] inner, [i] [0] * len(inner)) # 第 3 步声明本次迭代的向量 tile shapeUB 分块尺寸。 # 首维为 1当前 batch其余维度为内层输出形状 pypto.set_vec_tile_shapes(1, *inner_out) # 第 4 步在 a_s 内再取一片 rshape 为 [1] inner_out # 偏移全部为 0因为 a_s 已经是当前 batch 行 r pypto.view(a_s, [1] inner_out, [0] * len([1] inner)) # 第 5 步把 r 装配写回到输出 out 的对应位置 # 输出偏移 [i, 0, 0, ...]即 batch 轴写回第 i 行 pypto.assemble(r, [i] [0] * (len(ol) - 1), out)逐接口语义拆解pypto.loop(batch, namebatch, unroll_list[1])按batch长度建立循环。name用于在编译产物/Profile 中标识循环unroll_list声明允许的展开粒度当前骨架用[1]即一次处理 1 个 batch 行。仓库中更复杂的实现会使用带偏移与展开长度的pypto.loop_unroll例如 sum_lstm.py 中for bs_offset, unroll_length in pypto.loop_unroll( 0, batch_size, 1, nameLSTM_BATCH_LOOP, idx_namebs_offset, unroll_listtile_config.unroll_list ):可见unroll_list是可调优参数——sum_lstm.py 中 Arctic LSTM 将其配置为[1, 2, 4]以配合不同的 batch tile 大小。pypto.view(a, shape, offsets)从 GM 张量上按shape与offsets截取一个 view 切片不拷贝数据、不产生 GM 搬运。注意 offsets 列表的长度必须等于张量维数每个元素表示对应维度的起始下标。骨架第 2 步中[i] [0] * len(inner)表示「batch 轴偏移 i、其余轴从 0 开始整块取」。pypto.set_vec_tile_shapes(1, *inner_out)声明本次迭代的 Vector 单元 tile 形状UB 分块。首维固定 1单 batch 行后续维度来自inner_out。tile shape 决定 UB 装载粒度需按实际 shape / dtype / 平台约束调优而非照抄骨架。值得注意一个 kernel 内可以多次调用它来切换不同阶段的 tile 形状例如 sum_lstm.py 先用(current_tile_bs, h_tile)随后在门控阶段又切回(1, hidden_dim_4)。pypto.assemble(r, offsets, out)把片上计算结果r写回 GM 输出out的offsets位置。这里偏移[i] [0] * (len(ol) - 1)把第 i 行写回输出第 i 行其余轴从 0 开始——与pypto.view的读偏移形成「读一片 → 算一片 → 写一片」的对称结构。切分策略一句话总结骨架 Note 原文为batch 轴 loop 搬运沿切分轴取一片view 切片其余轴整块。这正是 PyPTO 轴切分算子的通用模式loop 覆盖 batchGM 搬运的单位view 负责在切分轴上精确定位切片其余轴整块随行搬运。占位符约定与最小可运行 setup骨架使用了一组占位符语义约定见 examples/README.md占位符含义sl输入 shape 列表如[B, S, D]ol输出 shape 列表pypto_dtype元素 dtype如pypto.DT_FP32batch被 loop 的外层轴长度通常sl[0]inner单次迭代处理的内层 shape如sl[1:]inner_out单次迭代的输出内层 shape末轴长度可能变化如 glu/diff/mean配套的最小可运行 setup同一个 README 提供import pypto B, D 8, 128 sl, ol [B, D], [B, 1] pypto_dtype pypto.DT_FP32 batch, inner, inner_out B, [D], [1]代入骨架后语义即为输入[8, 128]的 FP32 张量batch 轴8逐行 loop每行view出[1, 128]切片tile shape 为(1, 1)最终装配到输出[8, 1]。注意inner_out的末轴长度可以和输入不同如 split 出子序列、或经 glu/diff 等变换后维度变化这正是占位符拆分inner/inner_out的原因。真实源码佐证view assemble 的实战写法骨架是「参考形态」仓库中的生产级实现展示了更丰富的用法1. 多片切分同一输入多个 offsetssum_lstm.py 将融合后的[BS, 4H]张量按 4 等分切出 LSTM 四门pre_f pypto.view(fused, [current_tile_bs, hidden_dim], [0, 0]) pre_i pypto.view(fused, [current_tile_bs, hidden_dim], [0, hidden_dim * 1]) pre_o pypto.view(fused, [current_tile_bs, hidden_dim], [0, hidden_dim * 2]) pre_c pypto.view(fused, [current_tile_bs, hidden_dim], [0, hidden_dim * 3])这正是 split 的典型应用切分轴偏移 片序号 × 每片长度其余轴偏移 0。多个 view 共享同一份 GM 数据零拷贝切出多片。2. 带 valid_shape 的尾块处理gdr_bwd_impl.py 展示了 view 的第 4 个参数valid_shape用于动态长度如序列尾块act bt时声明有效形状pypto.view(q, [bt, head_dim], [off, hc], valid_shape[act, head_dim]), ... b_f pypto.fillpad(pypto.view(beta, [bt, 1], [off, h1], valid_shape[act, 1]), constant, 0.0)当 split 的切分块长度不能整除原轴长时尾片不足可组合valid_shape与pypto.fillpad做补零对齐该组合在 Qwen3-Next gated_delta_rule、chunk_kda 中用于 q/k/v/beta 尾块填充见 pypto-specific-ops.md 中fillpad条目。3. assemble 的多种写回位置gdr_bwd_impl.py 展示了assemble同时支持「view 切片后写回」与「整块直接写回」两种形态并支持 per-head 偏移[off, hc]累加写回。与 chunk 骨架的关系同一模式的不同切分粒度chunk.md 与 split.md 的骨架完全同构loop view set_vec_tile_shapes assemble 四段式差异仅在语义层split(split_size)按固定长度切分每片长度相同除最后一片可能不足chunk(chunks)按固定份数切分每片长度 ceil(轴长 / 份数)。无论哪种语义落到 PyPTO 上都是「先算各片 offsets再逐片 view assemble」这正是映射手册将二者归为同一条目split/chunk→view的原因。实际开发中每片对应的inner/inner_out与 tile shape 可能因片长不同而需要分别配置。使用约束与调优注意点骨架未经 NPU 编译验证examples/README.md 明确声明「骨架未逐一经 NPU 编译验证」且仅作「接口组合与轴切分模式」参考不是标准模板。loop 轴、unroll_list、tile shape、动态轴处理都需按实际 shape / dtype 与平台约束确定并调优。动态 shape 约束按 SKILL.md 的硬约束速查matmul / 归约类等计算 API 在编译期需要 concrete shape不接受含DYNAMIC维度的 tensor。若 split 算子中 batch 轴为动态如 sum_lstm.py 中的pypto.DYNAMIC应保证只有 loop 轴是动态的、tile 内维度为静态pypto.STATIC否则需采用「loop 切 tile」策略规避。tile shape 与 UB 容量set_vec_tile_shapes的取值需保证单次迭代数据能放入 UB且每维 0、最多 4 维不匹配会导致编译失败或性能劣化。可参考execution-constraints.md与 Tiling 文档pypto-set_vec_tile_shapes.md核对约束。view 是切片不是拷贝切出的 view 与原张量共享底层 GM 数据因此读侧不会产生额外搬运而写回必须通过assemble显式完成且输出 offsets 必须与循环变量对齐避免多核/多迭代写同一位置产生数据竞争。多片输出扩展骨架只展示了切出一片单输出的形态当需要把输入切成 N 片写到 N 个输出时在循环体内对每片重复「view → 计算 → assemble」即可各片 offsets 按片号 × 片长递增计算。小结PyPTO 中实现split的标准姿势可以概括为一条公式batch轴pypto.loop搬运 pypto.view按计算好的 offsets 取切片 pypto.set_vec_tile_shapes声明 UB 分块 pypto.assemble写回对应输出位置。仓库提供的 split.md 骨架是这一模式的最小化抽象而 sum_lstm.py多片切分、gdr_bwd_impl.py动态 valid_shape 尾块则展示了其在真实 NPU 算子中的完整演进形态——理解这四步即可举一反三地编写任意的轴切分类 kernel。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考