NPU DeepSeek-V3.2-Exp PyPTO 融合算子开发优化实践【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer简介近年 AI 大模型规模出现爆炸式增长其对算子性能的要求也越来越高。受限于硬件内存大小和带宽约束大模型时代的算子需充分考虑计算过程优化、内存读写优化、不同层级的内存复用等多方面因素而提升算子性能的重要手段之一便是算子融合。所谓算子融合是指通过将多个计算合并到一个算子 Kernel 中以减少内核启动开销和内存读写次数提高内存复用率从而提高算子性能。虽然算子融合性能收益显著但开发融合算子也常常伴随巨大的挑战融合意味着单算子中有多个计算逻辑如 MatMul、Softmax、ReLU代码量大复杂度高融合算子通常有多个输入输出对内存访问的要求高一旦核内不同层级的 SRAM 利用不善性能会大打折扣为了达到高性能需要使用大量当代硬件定制特性导致算子难以跨平台复用通常存在大量可配置参数手动调优困难且不可持续。随着融合算子的规模逐渐增大由于开发难度的原因手写融合算子的方式越来越难以为继。基于此算子编程领域兴起了一种针对现代硬件的高效易用计算编程范式——基于 Tile 编程。它将大规模计算任务和数据分解为更小、更易管理的“块”Tile从而更好地利用硬件的内存层次结构和并行计算能力。PyPTO 是 CANN 推出的一款高效的自研编程框架旨在简化算子开发流程大幅提升算子编程易用性同时保持高性能计算能力给开发者提供面向 MPMD 类似 Mega Kernel 开发的、基于 Tensor/Tile 的编程范式。相比于现有算子开发工具其主要特点包括自动生成读写指令开发者无需手动管理核内 SRAM。同时开发者可以通过参数配置控制 Tiling 切分和 SRAM 利用策略从而达到对核内 SRAM 的高效利用。定义了一组抽象的 Parallel Tile Operation (PTO) Instruction PTO 指令开发者编写的程序都会转换为一组 PTO 指令。通过在不同硬件平台实现对应 PTO 指令可自动完成算子的跨平台迁移不再需要手动适配。提供了一套 MPMD (Multiple Program Multiple Data) 运行时。相比于 SPMD (Single Program Multiple Data)MPMD 无需进行全局同步而是利用数据的生产消费关系构建任务间依赖使不同核在同一时间执行不同程序充分利用多核并行能力。提供了一套集成开发环境用于可视化算子编译及运行过程帮助开发者定位和调试算子精度/性能问题实现算子快速迭代。下图展示了 PyPTO 的主要架构PyPTO 架构特性介绍基于 Tensor/Tile 的编程范式PyPTO 是一个基于 Tensor/Tile 的编程范式其核心思想是使用 Tensor/Tile 作为数据的基本表达方式通过一系列对 Tensor 的基本运算来描述并组装完整的计算流程。Tensor 作为基础数据结构计算过程是对 Tensor 类型的数据进行数学运算。将原始 Tensor 进行切分后得到的小块数据称为 Tile而这个切分过程称为 Tiling。Tiling 使得处理器核可以根据核内 SRAM 容量大小容纳计算过程用到的 Tile 数据并复用这些 Tile 以达到减少对 Global Memory 访问的目的。基于 PyPTO 开发的代码整体风格简洁明晰开发者可以聚焦于算法本身的计算逻辑和各个 OP 的 Tile Shape 设置无需感知复杂的内存、调度、硬件等底层实现细节对于开发者更加友好开发难度大幅下降。Parallel Tile Operation在 PyPTO 中对数据的操作被称为 Operation。Tensor Operation 在经过 Tiling 之后转换为 Tile Operation。框架提供了一组基本的 Tile Operation 用于底层硬件抽象可针对不同芯片适配实现算子的跨平台支持。此外开发者可以自定义扩充 Tile Operation以支持新的硬件平台或者得到特定平台的高效性能。下表列举了一些 PTO 指令实例分类指令功能描述Element Op逐元素操作TADD执行两个 Tile 数据块内逐元素的加法运算......Tile-Scalar Op块与标量操作TADDS将标量与 Tile 数据块内逐元素相加......Axis Op按轴操作TROWSUM对 Tile 数据块的每一行执行求和......MATMUL Op矩阵乘操作MAMULB执行两个 Tile 矩阵数据块的乘法......Fix Pipe Op固定管线操作TMOV将 Tile 元素从源寄存器搬移到目的寄存器可带随路运算......MEM Op访存操作TLOAD将数据从内存加载到 Tile 数据块TSTORE将 Tile 数据块写回到内存......Complex Op复杂操作TCVT执行 Tile 数据块的类型转换......核内 SRAM 的管理PyPTO 为开发者提供两种 Tiling 调优的方式设定 Tile Shape 参数框架按内置 Tiling 方法将 Tensor Operation 计算图展开成 Tile Operation 计算图开发者可以自定义 Tiling 方法可以实现更特化更高效的展开逻辑。通过 Tiling 得到 Tile Operation 粒度的计算图后框架会对计算图进行切分并为切分后的子图插入数据搬运指令以使用各层级的核内 SRAM。每个子图的输入输出都放在 Global Memory 上而子图内只会使用核内 SRAM 用于存放中间数据。PyPTO 的子图切分过程会考虑以下因素Operation 亲和性在同一类型的核上执行的 Operation 会被尽可能分配到同一子图。子图同构性框架会尽可能保证切分出的子图是可复用的以减少编译耗时降低 ICache Miss 率。子图并行度和重复数据搬运的权衡过大的子图会导致核空闲过小的子图会导致同一份数据向不同核重复搬运。开发者可以通过配置参数调整切分策略在并行度和 SRAM 利用率间取得平衡。下图是一个算子的 Tiling 和切分过程示例图中展示了A乘以常数后与B相加然后与W进行矩阵乘最后对尾轴进行累加的一个计算示例。其中AB的维度为64, 64W的维度为64, 32。Tensor 图示例MulSAdd和RowSum均为Vector操作设置其Tile Shape为32, 32Matmul为Cube操作不进行切分。经过Tiling之后得到如下Tile计算图。Tile 图示例针对A3的CV分离架构框架会将Matmul单独切成一个子图Matmul前的MulS Add切成一个子图Matmul之后的RowSum切成一个子图。其中MulS_Add子图会被调用四次RowSum子图会被调用两次每次处理不同的Tile块。Execution Graph 图示例MPMD 运行时经过上一阶段的切图会得到许多小的子图这是 PyPTO 的最小执行单元每个子图会编译成可执行文件生成对应任务被分发到不同的核上执行。PyPTO 的运行时是 MPMD 的整个运行过程可分为如下几步根据子图间的数据依赖创建任务间依赖调度 AICPU 根据任务的拓扑顺序下发无前置依赖的任务到空闲 AI 核上AI 核收到任务后立即异步执行任务执行完成后调度 AICPU 收到回调信息解除后续任务的依赖并继续尝试下发。Human-In-The-Loop 的调优思路从算子描述到上板执行整个过程需要经过 Tile 展开、切图、运行时调度等众多阶段其中许多阶段面对的都是 NP-Hard 问题因此期望由框架自动生成高效性能的融合算子是不现实的。为此框架提供了一套可视化的集成开发环境使开发者可以查看每个阶段的输出信息控制各阶段的运行过程并将修改后的算子重新提交到 NPU 上执行得到新的结果和性能数据。通过这一机制开发者可以快速迭代在较短的周期内开发出性能优异的算子。下图展示了 PyPTO 的算子开发和调优过程PyPTO 工作流融合算子示例PyPTO 是 MPMD 的调度思想和 Mega Kernel 思想类似并支持开发者自行选定融合范围。本文以 DeepSeek-V3.2-Exp 中 Indexer Prolog、Deepseek Indexer Attention 和 Lightning Indexer 的开发为实例从融合算子开发者的视角阐述在 PyPTO 的编程框架下如何进行不同融合范围的算子的开发、以及通过 Human-In-The-Loop 的方式调优。Indexer PrologIndexer Prolog 是 DeepSeek-V3.2-Exp 的 Deepseek Indexer Attention 计算的前半部分其对应的 PyPTO 源代码可以参考PyPTO Deepseek Lightning Indexer Prolog其计算流程图如下Indexer PrologIndexer Prolog 的量化策略为Q_b_proj 使用 W8A8 量化其他 Linear 均不量化Indexer Q 使用 A8 量化Indexer Cache 使用 C8 量化反量化因子以 FP16 存储同时在量化前会进行 Hadamard 变换Hadamard 变换在业界最新的低 bit 量化方案中十分常用可以有效消除 outlier。当前该融合算子在典型 Shape4 BatchMTP1KV-Cache 64K 长度下kernel 性能已逼近理论的极限性能。计算公式Indexer Q 的计算公式如下$$ \bold{q}, \bold{q}{scale} \text{DynamicQuant}(\text{Hadamard}(\text{RoPE}(\text{DeQuant}(\bold{q} \cdot \bold{w}{qb})))) $$Q 的计算采用了动态的 Per-Token-Head 量化其中 Hadamard 变换通过矩阵右乘 hadamard_q 实现。而 $\bold{q}, \bold{w}_{qb}$ 均是 Int8 类型。Indexer Cache 的计算公式如下$$ \bold{k}, \bold{k}_{scale} \text{DynamicQuant}(\text{Hadamard}(\text{RoPE}(\text{LayerNorm}(\bold{x} \cdot \bold{w}_k)))) $$Cache 的计算同样采用了动态的 Per-Token-Head 量化其中 Hadamard 变换通过矩阵右乘 hadamard_k 实现。Indexer Weight 的计算公式如下$$ \bold{weight} (\bold{x} \cdot \bold{w}_{proj}) * \text{scale} $$Weight 的计算没有采用量化同时需要最后转化为 FP16 数据类型供后续的 Lightning Indexer 计算使用。泳道图与性能分析当前该融合算子在典型 Shape4 BatchMTP1KV-Cache 64K 长度下的性能 Profiling 数据即泳道图如下图所示Indexer Prolog Swimlane计算流拆解Indexer Prolog 的计算流呈现如下特点该计算包括了 Indexer Q, Indexer Cache 和 Indexer Weight 三个部分三个部分互相独立且每个部分串行执行。三条独立的计算流中Indexer Q 的耗时最长可以掩盖 Indexer Cache 和 Indexer Weight。在典型 Shape 下计算量较小并不会打满所有核进行计算性能瓶颈在搬运上。理论性能分析通过计算流拆解可知该融合算子的性能由 Indexer Q 的耗时决定其性能拆解如下第一部分是 (8, 1536) (1536, 8192) 的 Matmul 计算该计算是 MTE bound其理论的极限性能是 13us但由于同时会并行 k 和 weight 的 Matmul 计算抢占带宽从而可以说已逼近极致性能。第二部分是 DeQuant反量化和 RoPE 的计算。该部分计算均在 Vector 核上。通过合适的 tile 切分所有的计算均在 UB 上执行不会产生 UB 和 GM 之间的冗余搬运已达到极致性能。第三部分是 Hadamard 变换通过 (8, 64, 128) (1, 128, 128) 的 BatchMatmul 实现该计算仍然是 MTE bound已达到极致性能。第四部分是 Int8 量化计算该计算均在 Vector 核上。通过合适的 tile 切分所有的计算均在 UB 上执行已达到极致性能。综上所述Indexer Prolog 的计算已经逼近理论极限性能。下一步计划仿照 Indexer Prolog 的调优过程将其他几个子算子进行独立性能调优达到极致性能。Deepseek Indexer AttentionLightning Indexer 算子展示了 PyPTO 允许开发者可以控制算子融合的范围。接下来本文通过 Deepseek Indexer Attention 的融合示例展示 PyPTO 的整网融合能力。Deepseek Indexer AttentionDeepseek Indexer Attention 的计算逻辑图如图所示主要由 MLA Prolog、Indexer Prolog、Lightning Indexer、Sparse Flash Attention 4 个相对独立的部分构成可以看出在更大的范围内融合有更大的并行度可以挖掘。PyPTO 允许开发者先手动写这独立的四个部分并通过框架将四个部分融合成一个完整的 Deepseek Indexer Attention 大算子。因此大算子的开发策略如下对于复杂融合算子可以将复杂融合算子拆分成若干个相对独立的子算子根据泳道图对各个子算子进行独立性能调优。当各子算子接近最优性能后再把子算子们进行更大程度的融合以期获得完整融合算子的最优性能。其对应的 PyPTO 源代码可以参考PyPTO Deepseek Indexer Attention。与 Lightning Indexer 的融合算子相同也可以上板得到 Deepseek Indexer Attention 的泳道图DeepseekIndexerAttentionSwimlane泳道图分析大融合算子思想的优势主要在于将所有相对独立的子图融合在一起提高内存利用率并充分发挥NPU的性能。因此可以从图中看到DIA 整图会比单独的Lighting Indexer图更加复杂所有的计算步骤融合在一起乱序调度。从图中可以很清晰地快速定位到性能瓶颈点离散 Gather 部分的耗时占比高达 70% 以上。该部分的主要操作是 SparseFlashAttention 中根据 Lightning Indexer 计算得到的 TopK index从 kv_cache/kr_cache 中离散地读取对应 index 的 kv 进行聚合的过程。该聚合过程当前通过在 Vector 核搬运完成离散到连续的转换再写到 GM Workspace中由 Cube 核搬运到 L1 使用既浪费了内存又因重复搬运带来了极大的性能损耗。Indexer Prolog 子图 matmul 的任务过多。matmul M轴当前为1未进行合轴处理利用率过低。存在较大的流水线 bubble。两个子图之间存在未消除的barrier导致所有后续无依赖的任务也同样停滞出现了较大的 Pipeline bubble。reshape 空子图导致任务启动延时。reshape未充分优化导致存在的空子图后续依赖该 reshape 结果的任务停滞。改进措施及结果针对上述识别的关键性能瓶颈点逐一分析、解决PyPTO 允许开发者自定义 Tile Operation因此针对该场景需要开发特定的 TileOp其功能是根据 index 直接从 GM 离散搬运到 L1 的连续空间上使其能够直接进行 Matmul 的计算。这样可以省去离散转连续的额外workspace的申请并去除冗余的搬运性能提升非常明显。可以参考如下示意图自定义 Gather Tile Op将 b * s1 合轴调整 Tiling撑大 M 轴用满 L0A充分提高 Matmul 的计算效率性能倍增。消除两个子图之间的 barrier将两个子图充分融合不依赖的 task 可以提前启动可以消除流水线的气泡。reshape操作不改变数据排布消除后原本依赖 reshape 的 task 的 delay 消除并减少解依赖的时间。在经过上述分析和优化后得到新版的泳道图如下可以看到使能以上优化措施后该融合算子的运行时长缩短了 75% 。DeepseekIndexerAttentionSwimlaneNew下一步计划经过上述优化泳道图中的几个关键性能瓶颈点已经被消除性能也有了比较大的提升。基于新的泳道图我们可以继续挖掘性能优化点持续提升性能。Task 排布相对离散任务调度不够紧密需要分析任务之间的依赖消除 Pipeline Bubble 。SparseAttention 中离散的数据搬运效率较低需要提高搬运的并行度。多个子图之间仍然存在较大的空隙需要再进一步的融合将有数据依赖的计算整合到同一个 LOOP 中。期望经过以上优化后泳道图中的任务会排布更加紧密。总的来说性能优化是需要分多轮迭代的每一轮解决关键的性能瓶颈点后再基于新的泳道图进行分析、持续优化。Lightning IndexerLightning Indexer 是 DeepSeek-V3.2-Exp 的 Deepseek Indexer Attention 计算的一部分。其计算流程图如下LightningIndexer开发者根据计算流程图进行融合算子编码开发其对应的 PyPTO 源代码可以参考PyPTO Deepseek Lightning Indexer。基于 PyPTO 完成开发后 PyPTO 提供了完善的可视化集成开发环境为开发者提供了不同阶段的产物。其中最为重要的是在 NPU 上运行得到的性能 Profiling 数据即下图所示的泳道图开发者通过泳道图可以分析融合算子中各部分的耗时进而开展后续的性能调优。LightningIndexerSwimlane该泳道图的横轴表示时间纵轴由一条条的泳道构成一条泳道表示一个 AI 核。每条泳道上随时间排布了对应 AI 核上执行的任务。每个任务上标记了一个标签该标签是开发者写在源码中的语义标签 (Semantic Label)用来辅助开发者理解当前硬件在执行的任务对应到哪一段前端代码。选中一个任务时可以显示其前后依赖用来观察任务的执行顺序、依赖间隔等是否合理。在泳道图上可以看到不同颜色即前端开发者写的不同代码在同一时间不同计算核并行执行。传统的 SPMD 的调度要求同一时刻执行相同程序而 MPMD 的运行时会根据实际运行时解依赖情况来调度任务执行当一个任务解依赖完成后就可以和在不同 AI 核上的任务异步并行执行。泳道图分析当开发者写完 PyPTO 的代码获取第一版的开箱性能后需要根据泳道图分析可能的优化手段。上述泳道图展示的是在典型 Shape4 Batch KV-Cache 64K 长度下的执行情况直觉上能看到几个可以优化的“异常点”MatMul 的 Cube 任务前序无依赖为什么没有在第一时间执行主体的 Cube 和 Vector 任务都有三列但是第三列只有部分核在工作没有很好的对齐尾部。最后的 TopK 似乎有严重的拖尾。接下来PyPTO 的开发者可以针对上述问题进行分析。通过对异常点 1 的分析发现框架在处理数据搬运类 Operation 的时候对不需要数据搬运的场景优化能力缺失。通过对异常点 2 的分析发现由于是 4 Batch 64K 的场景Page Attention 中给定的 Block Size 为 128那么 MatMul 切分出了4*64*1024/1282048个任务程序中通过配置调整任务 L1 复用为 32即最后有2048/3264个任务无法平均分配到各个核上当前硬件共有 24 个 Cube 核。同样由于是 4 Batch 的场景且最终的 TopK 任务是单核的归并所以在 Vector 核上有拖尾。我们可以尝试优化最后的排序、归并任务来减小拖尾。改进措施及结果针对异常点 1 通过修改相应的 Pass 补齐泛化能力从而提升性能 (10%) 。与传统黑盒的调优框架不同 PyPTO 通过可视化的集成开发环境辅助开发者发现算子甚至框架的问题开发者可以选择本地修复或者向开源社区反馈。针对异常点 2 开发者可以选择通过调整 L1 复用来达到特定 Shape 性能和动态 Shape 泛化性能的平衡。只是在本例中L1 复用设置为 32 可以达到较好的泛化性。config::SetPassOption(L1_REUSE, 32);在对异常点 2 分析的过程中发现 MatMul 的输出粒度为 128进而发现 Tile Shape 同样设置为 128这对硬件计算并不是最高效的因此选择调整 Tile Shape让 MatMul 把数据输出到连续内存后再进行后续大颗粒度的 Vector 运算提升性能 (17%) 。TileShape::Current().SetVecTile({64, 256});在经过上述优化后得到新版的泳道图如下可以看到其性能有明显的优化 (27%) 。LightningIndexerSwimlaneNew下一步计划经过上述优化泳道图已经比较密集。但是图中仍然有一些可优化的部分MatMul 的子图太大会导致后续 Vector 任务延迟开始第一列 Cube 任务下方仍然有大量的空白时间。通过对原始计算逻辑的分析得知这一部分可以并行因此后续可以考虑通过控制复用程度对此进行优化。最后的 TopK 是通过排序归并得出的结果可以考虑将排序部分放入前面 Vector 运算最后只做归并从而减小拖尾时间。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
