CANN/DeepSeek-V4 PyPTO融合算子优化

NPU DeepSeek-V4模型 PyPTO 融合算子优化【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer面向 DeepSeek-V4架构本次发布多个全新的融合算子hc_preCompressor和三个 FA 类算子SWA,CFA和SCFAMlaProlog和IndexerProlog也针对新架构进行了适配。其中hc_pre是本次 DeepSeek 新架构中Manifold-Constrained Hyper-Connections(mHC) 的一部分。本篇技术报告中我们将对这些融合算子在 PyPTO 上的具体的实现和优化进行概述。同时我们也会介绍 PTO-ISAPTO 虚拟指令集帮助 PyPTO 实现跨代际兼容。若对 PyPTO 高性能编程框架代码感兴趣可参考 PyPTO 仓库Cann-PyPTO。以下是单 die 64 batch, MTP1序列长度 8k 下各算子的实测性能数据以及各算子的代码规模可以看出 PyPTO 可以以较少的代码量高效的开发效率3-7天完成复杂融合算子的开发且能取得不错的基础性能算子代码行数(KLOC)BS1S2Dcmp_ratiocoffstart posDuration(us)sliding_window_attention0.36428192512N/AN/AN/A65compressed_flash_attention0.26428192512128N/AN/A90sparse_compressed_flash_attention0.264281925124N/AN/A160compressor_c4li0.9642N/A12842819255compressor_c4a642N/A51242819280compressor_c128a642N/A5121281819255compressor_c4li642N/A12842819175compressor_c4a642N/A51242819195compressor_c128a642N/A51212818191125mla_prolog0.5642N/AN/AN/AN/AN/A130hc_pre0.4642N/AN/AN/AN/AN/A70lightning_indexer_prolog0.3642N/AN/AN/AN/AN/A101hc_prehc_pre在每层会调用两次分别在 attention 和 ffn 的计算前。计算过程计算 RMSNorm 的分母$$ rsqrt \sqrt{\frac{1}{\frac{1}{n}\sum_{i1}^n x_i^2 \epsilon}} $$计算 mixes$$ mixes (x hc_fn) \odot rsqrt $$Sinkhorn-Knopp 算法$$ pre, post, comb sinkhorn(mixes, hc_scale, hc_base, hc_mult, hc_sinkhorn_iters) $$Sinkhorn-Knopp 算法每次迭代会进行逐行归一化再做逐列归一化hc_sinkhorn_iters 控制迭代次数。利用 pre 和 x 计算 y$$ y rowsum(pre \odot x) $$实际开发中需要充分利用昇腾硬件的特性以实现更好的性能以下将详细介绍当前的解决方案及具体的实现方式。流水并行计算流程中第一步是 Vector 计算而第二步中 $x hc_fn$ 是矩阵乘的 Cube 计算且两者无依赖关系。PyPTO 的框架支持了两个任务的同时执行无需手动插入同步。指令优化在 Sinkhorn-Knopp 算法中涉及到了大量的重复的归约串联逐元素乘除的计算在逐元素计算中通过 pass 优化和 PTO 指令避免了显式地广播实现了高性能的计算。pypto.experimental.set_operation_config(combine_axisTrue)MatMul优化decode场景下matmul计算的K轴很长16K而M\N轴较小因此采用切K的方式将数据切分到不同的核上以充分利用计算资源pypto.set_cube_tile_shapes([16, 16], [512, 2*1024], [128, 128], \ enable_multi_data_loadTrue, enable_split_kTrue) mm_res pypto.matmul(x_fp32, hc_fn, pypto.DT_FP32, b_transTrue)计算优化Prefill 和 Decode 的最大区别在 x 入参的的 0 轴大小上。Prefill 场景下t 通常比较大而在 Sinkhorn-Knopp 算法中归约轴很小为 hc_mult4且在行归约时为 -1 轴列归约时为 -2 轴。此时我们采用数学上等价计算的技巧在第二步计算中在矩阵乘中直接计算转置$$ hc_fn^T x^T $$从而在进入 Sinkhorn-Knopp 算法中t 会置于尾轴上归约轴为 -2, -3 轴这样能提高计算的效率。泳道图在单 die 64 batch, MTP1 的场景下上板执行结果如下hc_pre泳道图MlaProlog该算子是 Multi-Head Latent Attention 将 hidden states 处理为 query 和 kv 的计算操作。计算过程Query 的计算包括两次采样和 RmsNorm其中第二次 RmsNorm 权重恒为 1最后对 -1 轴的后 rope_dim 维度进行 inplace interleaved rope 计算$$ c^Q RmsNorm(x wq_a) $$$$ q RmsNorm(c^Q wq_b) $$$$ q[..., -rope_dim:] ROPE(q[..., -rope_dim:]) $$Kv 的计算包括一次采样和 RmsNorm最后对 -1 轴的后 rope_dim 维度进行 inplace interleaved rope 计算$$ k RmsNorm(x wkv) $$$$ k[..., -rope_dim:] ROPE(k[..., -rope_dim:]) $$流水并行PyPTO 采用了 MPMD 的调度方式在 Decode 场景下像xwq_a和xwkv不会使用满所有的核因此可以并行执行同理query 和 kv 后续的 RmsNorm亦可并行执行Unrollist 将动态展开为静态图在 Prefill 和 Decode 场景中t 的差异往往是巨大的。通常情况下一套代码Tile 切分很难在不同场景下均达到较好的性能但是通过对动态值 t 进行了多分档以及 Tile 切分采用静态 tTile 的方法达到了在不同场景下均实现较优性能的结果。unroll_list [128, 64, 32, 16, 1] for tIdx, unrollLength in pypto.loop_unroll(0, t, 1, nameMLA_BS_LOOP, idx_namebs_offset, unroll_listunroll_list, ): ...合图优化通过设置合图参数可以将相同结构的子图合并避免同一结构子图数过大、避免重复搬运同时可以增大核内流水调度可能性实现了较优的性能。pass_options{ cube_nbuffer_mode: 1, vec_nbuffer_mode: 2, vec_nbuffer_setting: {-1: 2} }MatMul优化根据输入参与matmul计算的tensor的shape通过设置合适的tile_shape将matmul的任务均衡地分配到各aic 核上。当k轴过大时使能切k或者手动切k减少了单个任务的搬运量。在shape较大的情况下开启大包搬运提高了带宽利用率。通过以上方式提升了matmul任务的效率。# x_tile.shape[128, 4096], wkv_shape[4096, 512] pypto.set_cube_tile_shapes([tile_bs, tile_bs], [256, 512], [128, 128], enable_split_kTrue, enable_multi_data_loadTrue) kv pypto.matmul(x_tile, wkv, pypto.DataType.DT_FP32)泳道图在单 die 64 batch, MTP1 的场景下上板执行结果如下mla_prolog泳道图Compressor新模型引入了Compressor模块PyPTO 对此进行了实现。计算过程$$kv XW^{KV}$$ $$score XW^{Gate}ape$$ 使用kv和score对kv_state和score_state进行刷新。推理场景下当start_pos1为compress_ratio倍数时按下面的计算生成压缩的kv $$kv(kv_state * score_state.\text{softmax(dim1)).sum(dim1)}$$ $$kv\text{ROPE}(\text{RMSNorm}(kv))$$rotatetrue时对kv进行Hadamard变换。多场景的不同优化Compressor模块使用的参数例如压缩比compress_ratiorotate等配置在各层之间有所不同。针对不同的场景配置PyPTO 在同一个算子中根据不同的分支在每个分支上进行各自的优化从而在不同的使用场景下均能达到较为良好的性能。下面以compress_ratio4rotatefalse时的场景为例介绍优化方法。MatMul优化b 64 b_loop (bsz b - 1) // b for b_idx in pypto.loop(b_loop, nameLOOP_COMP_1, idx_nameb_idx): x_view pypto.view(x, [b*s1, h], [b_idx*b*s1, 0]) ## Matmul pypto.set_cube_tile_shapes([128, 128], [256, 512], [128, 128], True) kv_t pypto.matmul(x_view, wkv, pypto.DT_FP32, b_transTrue) score_t pypto.matmul(x_view, wgate, pypto.DT_FP32, b_transTrue)该场景下两个矩阵乘的左右矩阵shape为[batchsize*s1, 4096], [4096, 1024]。batchsize是动态的。如果直接对batch进行循环展开静态图中矩阵乘的M轴shape将为s1导致整个算子中右矩阵被重复搬运多次性能较差。因此在动态的batch轴我们使用b64进行切分。K轴使能大包搬运cube_tile_shapes([128, 128], [256, 512], [128, 128], True)代表K轴的L0和L1搬运的tilesize分别设置为256, 512。通过一次搬运较大的数据量进入L1来提升带宽使用率。向量操作优化b_valid (bsz-b_idx*b).min(b) for c_idx in pypto.loop(b_valid, nameLOOP_COMP_2, idx_namec_idx): pypto.set_pass_options(pg_skip_partitionTrue)由于不同batch的start_pos不同它们的vector处理可能进入不同的分支。不同batch所更新的kv_state和score_state的地址也不同。因此我们对batch进行loop对每个batch分别进行vector处理。我们通过设置pg_skip_partitionTrue让一个batch的vector处理都在一个子图内。这是因为断图时中间数据会在GM和UB之间搬入搬出造成性能劣化。vector处理包括SoftmaxReduceSum和RMSNormRoPE前者在1轴上进行操作后者在尾轴上进行操作难以较好地进行多核并行。在高吞吐的大batchsize场景时由于batch间是独立的操作因此每个batch作为一个子图可以多核并行达成较好的性能。index cache_index[:,pos:pos1] kv_state scatter_update_3d(kv_state, index, kv) ## b,4,2d score_state scatter_update_3d(score_state, index, score)对后续vector操作使用的kv_state和score_state更新时使用scatter_update这一OP可以避免将kv_state和score_state多次搬入UB以及concat操作导致的性能劣化。泳道图在单 die 64 batchMTP1并且每个batch都存在压缩操作的场景中上板结果如下compressor泳道图IndexerProlog和 DeepSeekV3.2 不同的是IndexerProlog 只会出现在 compress_ratio 4 的层为后续 LightningIndexer 计算提供输入q、weight 及 q_scale。计算过程Query 的计算 q, q_scale的计算公式如下$$ q_tmp \text{qr}{idx_wq_b} \cdot \text{qr_scale} \cdot \text{idx_wq_b_scale} $$$$ q_hadamard \text{Cat}({q_tmp[:, :nope_dim], Rope(q_tmp[:, nope_dim:])}, -1)hadamard $$$$ q, q_scale Quant(q_hadamard) $$其中Rope表示旋转位置编码计算Quant表示量化计算。Weights 的计算 Weights的计算公式如下$$ weights x\text{weights_proj} \cdot {\frac{1}{\sqrt{\text{idx_nq} \cdot \text{head_dim}}}} $$泳道图在单 die 64 batchMTP1 的场景下上板结果如下indexer_prolog泳道图FA 类算子SWA、CFA 和 SCFA本次新模型的一个显著特点就是各层之间根据压缩比的不同会采用不同的 Attention 计算方式。PyPTO 充分发挥了开发灵活性的特点将不同层的 attention 实现分为三个不同的 attention 算子进行开发SWASWA是 sliding window attention 的缩写即滑窗注意力。在 compress_ratio 1 的场景下会采用该算子。通过掩码实现了在 Prefill 阶段和不同 MTP 配置下的高效实现。泳道图在单 die 64 batchMTP1kv_len8k 的场景下上板结果如下swa泳道图CFACFA是 compressed flash attention 的缩写。在 compress_ratio 128 的场景下会采用该算子。注意到CFA 的 kv cache 有两个来源分别是 ori_kv 和 cmp_kv其中 ori_kv 是原始的 kv cachecmp_kv 是经过compressor模块压缩后的 kv_cache。Ori_kv 我们采用滑窗选取而 cmp_kv 需要满足因果注意力机制。泳道图在单 die 64 batchMTPkv_len8k 的场景下上板结果如下cfa泳道图SCFASCFA是 sparse compressed flash attention 的缩写。在 compress_ratio 4 的场景下会采用该算子。和 CFA 不同的是SCFA 中对 cmp_kv 会通过 cmp_sparse_indices 离散地选取指定的 kv cache。我们通过 UB 将离散的 token 进行连续的聚合使后续的 Attention 计算高效。泳道图在单 die 64 batchMTP1kv_len8k 的场景下上板结果如下scfa泳道图PTO-ISA 支持不同代际芯片的算子兼容在算子层面除了像 MXFP8 量化这类由于不同代际芯片本身能力不支持外算子前端无需感知芯片的代际差异即一套算子代码即可在不同代际的芯片上运行。而芯片不同代际的能力通过 pass IR 和后续的 PTO 指令进行区分。PTO-ISAPyPTO 后端基于 PTO 虚拟指令集PTO ISA打造。 PTO ISA 兼容新旧核架构A2/A3/950D并做到“兼容不降速”针对不同代际芯片统一指令接口并做到微架构高度适配。 性能方面封装 DSA 定制优化结合毕昇 PTO 编译器强大的 VF 融合能力更好进行micro kernel级别的融合充分发挥新一代芯片算力。同时PTO ISA及配套毕昇编译器面向算法快速更新的开发需求提供了强大的 DFX 能力提供 TPRINT (打印 Tile 或 GlobalTensor)、算子融合信息统计等功能方便精度调试与软件优化。hc_posthc_post 的计算如下y post.unsqueeze(-1) * x.unsqueeze(-2) torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim2)只会使用到 Vector 核。在已做其他优化基础上PTO-ISA 结合毕昇 PTO 编译器强大的 VF 融合能力通过自动融合性能相比于未使能自动融合提升22.7%。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考