ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

PyPTO-Gym SK-09 向量多轴分块骨架:多维权重分块、动态展开与 Pipeline 重叠的 NPU 向量算子设计范式

PyPTO-Gym SK-09 向量多轴分块骨架:多维权重分块、动态展开与 Pipeline 重叠的 NPU 向量算子设计范式 PyPTO-Gym SK-09 向量多轴分块骨架多维权重分块、动态展开与 Pipeline 重叠的 NPU 向量算子设计范式【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym本文围绕 pypto-gym 算子设计模式库中的 SK-09「Vector Tiling (Multi-Axis)」骨架展开它规定了纯向量V算子在两个及以上维度上分块的循环结构、tile 加载方式与展开因子选择规则。读完本文你将掌握如何在 PyPTO 中设计 batch序列如 RoPE或参数 M×N如 AdamW多轴分块 kernel如何正确处理尾块以及优化器场景下submit_before_loopNONE_CACHEABLEpipeline 双开关的原理与配套调优参数并能在仓库中找到 InterleaveRope、ApplyAdamWV2 两个完整实证案例进行对照阅读。一、SK-09 的定位模式库中的纯 V 多轴分块骨架SK-09 定义在 SK-09-vector-multi-axis.md与 骨架索引 中其余 15 个骨架online flash attention、线性投影、MoE 等并列是 pypto-gym 算子设计流程pypto-op-design中面向纯向量计算的一类核心结构范式。其文档头明确标注CV 排布纯 V不涉及 Cube 矩阵乘典型算子InterleaveRope、ApplyAdamWV2、ApplyRMSProp适用场景需要在多个维度上分块的向量操作——例如 RoPE 在 batchseq 两个维度分块优化器在参数维度分块。它解决的核心问题是当输入张量的有效数据量远超片上可一次容纳的 tile 规模时如何用「外层循环遍历大轴 内层循环展开小轴」的组合把一个大向量算子切成一串可流水化提交的 tile 计算同时保证尾块不能整除的正确性与性能参数nbuffer、展开因子、cache 策略的配套。二、骨架结构多轴循环 view/valid_shape assemble原文档给出的最小骨架如下它完整体现了「外层普通循环、内层展开循环、逐 tile 加载、单 TileShape 计算、逐 tile 写回」五步结构def vector_tiling_kernel(input, cos, sin, output): for b_idx in pypto.loop(B): # Loop: Outer axis for s_idx, uf in pypto.loop_unroll(S, unroll_list[32]): x_tile pypto.view(input, [uf, N, D], [b_offset, 0, s_offset], valid_shape[uf, N, valid_s]) cos_tile pypto.view(cos, [uf, 1, D//2], [b_offset, 0, s_offset], valid_shape[uf, 1, valid_s]) sin_tile similar # V: Multi-Axis Vector Compute set_vec_tile_shapes(v_rows, v_cols) ... compute ... pypto.assemble(result, [b_offset, 0, s_offset, 0], output)逐行对照仓库实证代码可以确认几个关键约定外层pypto.loop承担大轴遍历batch / head / 参数 M 维每轮只处理一个「切片」内层pypto.loop_unroll承担小轴序列/参数 N 维的定长 tile 遍历展开因子通过unroll_list指定文档规定展开因子候选为32、16、1初始设计只选一个值且必须验证不能整除时的处理其余值作为后续调优候选pypto.viewvalid_shape完成逐 tile 加载valid_shape是尾块处理的统一入口pypto.assemble完成逐 tile 写回偏移与 view 的偏移严格一致保证写回位置与加载位置对齐。三、关键编码特征六条设计规则原文档归纳的编码特征表完整保留如下它是判断一个 kernel 是否属于 SK-09 骨架的 checklist特征规则多轴嵌套外层pypto.loopbatch/param 维度内层loop_unroll序列维度内层 unroll内层 loop 使用loop_unrolltile_size 动态变化view 加载多层嵌套场景用pypto.viewvalid_shape逐 tile 加载TileShape 固定全程一种 vec_tile 配置无切换Pipeline Loop优化器场景使用submit_before_loopTrueset_cache_policy(NONE_CACHEABLE)无跨 loop 状态各 tile 计算独立无累积器典型算子InterleaveRope, ApplyAdamWV2, ApplyRMSProp其中两条规则值得结合源码展开无跨 loop 状态意味着 SK-09 与需要 softmax 归一化、online max 的 attention 骨架SK-01/SK-02有本质区别每个 tile 的输入输出自洽这是它能安全做 pipeline overlap 的前提。若你的算子存在跨 tile 的归约累加器如 running mean应改用 SK-10 recurrent-state 一类的骨架而不是硬套 SK-09。TileShape 固定指向量计算块set_vec_tile_shapes在单条 tile 计算流内保持稳定。注意这并不禁止「先按 gather 域设置一次、再按元素域设置一次」的两段式切换——InterleaveRope 中正是这种 gather/element 两段 TileShape 的写法但每一段内部不随 tile 循环变化。四、实证案例一InterleaveRope 的三轴 batchheadseq 分块interleave_rope_impl.py 是 SK-09 在 RoPE 场景的完整落地4D 输入x: [B, 128, S, 64]B、S 为动态轴N∈{1,128}D 固定 64按b 轴外层 loop → n 轴中层 loop → s 轴内层 unroll的三轴结构分块。核心 kernelbf16/N128 主路径的循环与加载结构如下pass_options{vec_nbuffer_setting: {-2: 1, 1: 8}}, # jit 装饰器级 nbuffer 配置 for b_idx in pypto.loop(batch, nameb_loop): # 外层batch 轴 for n_blk in pypto.loop(128 // n_length, namen_loop): # 中层head 轴 n_off n_blk * n_length for s_blk, unroll_length in pypto.loop_unroll( # 内层seq 轴动态展开 0, seq_len, 1, names_loop, idx_namebs_blk_offset, unroll_listtile_config.unroll_list ): x_t pypto.view(x, [1, n_length, unroll_length, dim], [b_idx, n_off, s_blk, 0], valid_shape[1, n_length, unroll_length, dim]) c_lo pypto.view(cos, [1, 1, unroll_length, HALF], [b_idx, 0, s_blk, 0], valid_shape[...]) # ... 同一 loop 内继续 view c_hi / s_lo / s_hi ... pypto.set_vec_tile_shapes(gather_tile[0], gather_tile[1], gather_tile[2], gather_tile[3]) # gathermask 奇偶抽取 → cast fp32 → mul/sub/add → cast 回 bf16 pypto.assemble(ye, [b_idx, n_off, s_blk, 0], out) pypto.assemble(yo, [b_idx, n_off, s_blk, HALF], out)见 interleave_rope_impl.py#L64-L131。对照骨架规则可以验证展开因子动态化unroll_list{64, 32, 1}RopeTileConfig第 55-61 行loop_unroll在不能整除 seq_len 时以较小的 unroll_length 处理尾块——这正是文档要求「验证不能整除时的处理」的实证多 view 同步加载x、cos 前后半区、sin 前后半区共 5 个 view 全部在同一条 s_loop 内发出让编译器识别为同 batch 提交这是文档性能表中「多 view 同步加载」推荐项的直接来源尾块契约InterleaveRope README 明确「S 任意值含 SS_TILE 与不被 S_TILE 整除均支持通过 ceil-div valid_shape处理尾块」并约定输出为 split-half layoutout[..., 0:32]为偶位、out[..., 32:64]为奇位下游 attention 的 Q/K 必须同为 split-half 排布才能保持点积等价wrapper 按 (arch, N, dtype, S, S_cs) 派发 4 套 kernel 变体第 603-652 行950 架构走deinterleave单指令变体、S1/S2 走短序列变体、S_cs1 走 broadcast 提升变体cos/sin view 提升到 b 循环外复用。从源码结构看同一 SK-09 骨架可通过 wrapper 派发覆盖全 shape 域而不必在单 kernel 内做条件分支。对应精度测试与用例见 test_interleave_rope.py 与 test_cases.json容差为 atol1e-4、rtol7.8125e-3。五、实证案例二ApplyAdamWV2 的参数维 M×N 双轴分块优化器场景是 SK-09 的另一典型形态权重为 2D[M, N]沿参数 M 轴和列 N 轴分块。apply_adam_w_v2_impl.py 实现了带 bias correction 的单步 AdamW 更新m_t/v_t递归、m_hat/v_hat修正、解耦 weight decaykernel 主体结构for m_idx in pypto.loop(meta.m_loops, nameadamw_m_loop_bf16, unroll_list[1]): m_offset m_idx * meta.m_tile valid_m (meta.m_dim - m_offset).min(meta.m_tile) # 尾块 for n_idx in pypto.loop(meta.n_loops, nameadamw_n_loop_bf16, unroll_list[1]): n_offset n_idx * meta.n_tile valid_n (meta.n_dim - n_offset).min(meta.n_tile) valid_shape [valid_m, valid_n] w_tile pypto.view(weight, [meta.m_tile, meta.n_tile], [m_offset, n_offset], valid_shapevalid_shape) # 同 loop 内 view grad / m / v 四个输入 tile # cast fp32 → add/mul/div/sqrt 更新流 → 三路 assemble 写回 pypto.assemble(w_new_bf16, [m_offset, n_offset], weight_out) pypto.assemble(m_new, [m_offset, n_offset], m_out) pypto.assemble(v_new, [m_offset, n_offset], v_out)见 apply_adam_w_v2_impl.py#L164-L192。与骨架规则的对应点动态 tile 元信息_adam_tile_meta第 54-62 行用pypto.ceildiv计算两个轴的 loop 数tile_config中可携带第 4 个元素作为 N 轴 tile 大小——这正是文档中「K_TILE 大小2048优化器经验值」在通用化 kernel 里的体现TARGET_K_AXIS_TILE_CONFIG [7168, 16, 512, 4096]面向[7168, K]网络形状特化精度路由bf16 路径入口cast到 fp32、出口前 cast 回 bf16m/v 状态量恒为 fp32第 177-188 行step只通过 host 预计算的bc1/bc2标量进入 kerneldevice 侧不做pownbuffer 配置bf16 kernel 的 jit 装饰器配pass_options{vec_nbuffer_setting: {-1: 4, -2: 1}}第 133-139 行K 轴特化路径额外在 kernel 内set_pass_options(vec_nbuffer_setting{DEFAULT: 4})并开启combine_axisTrue与 InterleaveRope 的装饰器级配置互为参照。ApplyAdamWV2 README 给出了可复制的最小用法apply_adam_w_v2_wrapper(weight, grad, m, v, beta10.9, beta20.999, lr1e-3, weight_decay0.01, eps1e-8, step1)与精度测试入口tests/ops/experimental/vector/ApplyAdamWV2/test_apply_adam_w_v2.pylevel0 fp32、level1 bf16atol1e-4、rtol0.0078125通过后输出[PRECISION_PASS]。六、Pipeline Loop优化器场景的特化路径对于 AdamW 这类「大参数张量 × 逐 tile 独立计算」的算子瓶颈通常不在计算而在大张量加载阻塞计算。SK-09 为此提供了一条特化路径即原文档「Pipeline Loop」小节for k_idx in pypto.loop(k_loops, submit_before_loopTrue): # pipeline overlap tile view(param, [M, N_TILE], [0, k_off], valid_shape[M, valid_n]) set_cache_policy(NONE_CACHEABLE) # 输入不做缓存 ... compute ...两个开关必须理解为一组而不是两个独立的调优项submit_before_loopTrue把循环内各 tile 的加载/计算提交重叠起来形成 pipeline overlap前提是「无跨 loop 状态」这一骨架特征成立set_cache_policy(NONE_CACHEABLE)对大参数输入关闭 cache 一致性。文档明确指出单开submit_before_loop但保留默认 cache 策略会因 cache 冲突反而变慢——大参数张量把一致性 cache 打满后重叠带来的收益被一致性维护开销吃掉。从仓库结构看submit_before_loop在 chunked_gated_delta_rule_impl.py、kda_chunk_impl.py 等「向量计算 循环耦合」的 chunk 型算子中同样被使用原文档将 chunked_gdr 实现约 L411-423 区域列为「向量循环耦合」的实证来源之一另一来源是 mla_prolog_v4_impl.py 的 RoPE 多轴部分文档标注 L312-360可作为 SK-09 在 MLA 前处理链中的延伸阅读。七、开箱性能优化参数表逐项解读原文档附带的性能配置表是 SK-09 落地时的主要调优清单完整继承并结合两个实证 kernel 的取值展开如下维度推荐配置取值经验作用仓库实证pass_options.vec_nbuffer_setting必配{-1: 4}起步多轴重叠时升至{0: 4, -1: 6}多 view tile 并发加载需要更多 nbufferInterleaveRope 用{-2: 1, 1: 8}/{-2: 1, -1: 4}AdamW bf16 用{-1: 4, -2: 1}循环展开与 Vector TileShape配合评估每轮数据块大小与所选展开因子一致计算 TileShape 另按 API 设置16 可作为单值候选并验证尾块RoPEunroll_list{64,32,1}submit_before_loopTrue必配优化器场景内层 loop 标记启用 pipeline overlap把权重加载和计算重叠见第六节set_cache_policy(NONE_CACHEABLE)必配优化器场景大参数张量上避免大参数挤占 cache 一致性带宽见第六节多 view 同步加载推荐x_tile / cos_tile / sin_tile 同 loop 内 view让编译器识别为同 batch 提交InterleaveRope 5 view 同 loop内层 loop 嵌套深度限制≤ 3 层超过 3 层会触发 18000 表达式上限RoPE 恰好 3 层b/n/sK_TILE 大小推荐2048优化器/128RoPE参数维分块经验值AdamWN_TILE1024~4096梯度化combine_axisTrue推荐仅当存在尾轴-1 broadcast 二元运算时尾轴 broadcast 内联 brcb两个实证 kernel 均开启两条使用要点嵌套深度 ≤ 3 层是硬约束而非建议。InterleaveRope 的 b→n→s 三层循环已是上限若你的算子天然需要四轴遍历应把某个轴合并进 tile 维度如 batch 合并进 N_TILE或改用 wrapper 派发多个 kernel而不是继续加深嵌套。combine_axisTrue与尾轴 broadcast 绑定。AdamW 的 K 轴特化路径只在len(tile_config) 3时才开启它第 160-162 行从源码结构看这是把「存在尾轴-1 broadcast 二元运算」这一启用条件工程化到了配置里。该骨架特有的性能方向可概括为一句话submit_before_loop NONE_CACHEABLE pipeline 双开优化器特化。两个开关一起才能形成 pipeline overlap单开submit_before_loop但保留默认 cache 策略会因 cache 冲突反而变慢。八、边界判断什么时候选 SK-09把原文档的场景描述与两个实证案例合并SK-09 的适用边界可以归纳为四条判据纯向量计算flow pattern 仅含 V无 Cube——涉及 GEMM 的融合应看 SK-15 general-cv-fusion 一类骨架有效数据量超过单 tile 容量需要在 ≥2 个维度上分块tile 间无累积状态无 softmax 归一化、无 running 统计量每个 tile 可独立提交与流水化存在动态轴且需要尾块正确性——统一走 ceil-div loop 数 valid_shape收缩的方案而不是在 kernel 内做 if 分支。设计流程上SK-09 属于 pypto-op-design 的 skeleton 层先用 atom 层如 AT-04 rope、AT-19 optimizer-update分解计算原语再套入本骨架的循环/view/assemble 结构最后按第七节参数表完成 nbuffer、展开因子与 pipeline 开关的初始配置并用对应测试tests/ops/experimental/vector/下的 level 分级精度测试闭环验证。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表