CANN PyPTO 融合算子开发优化实践以 DeepSeek-V3.2-Exp Indexer 系列算子为例【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer本文基于 CANN 的 PyPTO 编程框架系统讲解基于 Tensor/Tile 的融合算子开发范式、核内 SRAM 管理、MPMD 运行时与 Human-In-The-Loop 调优方法并以 DeepSeek-V3.2-Exp 中 Indexer Prolog、Lightning Indexer、Deepseek Indexer Attention 三个融合算子的开发与调优为实战案例。读者读完后将掌握 PyPTO 的编程模型、如何通过泳道图定位性能瓶颈以及从单算子融合到整网融合的迭代优化方法论。背景大模型时代算子融合的挑战与 Tile 编程范式近年 AI 大模型规模出现爆炸式增长其对算子性能的要求也越来越高。受限于硬件内存大小和带宽约束大模型时代的算子需充分考虑计算过程优化、内存读写优化、不同层级的内存复用等多方面因素而提升算子性能的重要手段之一便是算子融合。所谓算子融合是指通过将多个计算合并到一个算子 Kernel 中以减少内核启动开销和内存读写次数提高内存复用率从而提高算子性能。虽然算子融合性能收益显著但开发融合算子也常常伴随巨大的挑战融合意味着单算子中有多个计算逻辑如 MatMul、Softmax、ReLU代码量大复杂度高融合算子通常有多个输入输出对内存访问的要求高一旦核内不同层级的 SRAM 利用不善性能会大打折扣为了达到高性能需要使用大量当代硬件定制特性导致算子难以跨平台复用通常存在大量可配置参数手动调优困难且不可持续。随着融合算子的规模逐渐增大由于开发难度的原因手写融合算子的方式越来越难以为继。基于此算子编程领域兴起了一种针对现代硬件的高效易用计算编程范式——基于 Tile 的编程。它将大规模计算任务和数据分解为更小、更易管理的块Tile从而更好地利用硬件的内存层次结构和并行计算能力。PyPTO 是 CANN 推出的一款高效的自研编程框架旨在简化算子开发流程大幅提升算子编程易用性同时保持高性能计算能力给开发者提供面向 MPMD类似 Mega Kernel开发的、基于 Tensor/Tile 的编程范式。相比于现有算子开发工具其主要特点包括自动生成读写指令开发者无需手动管理核内 SRAM。同时开发者可以通过参数配置控制 Tiling 切分和 SRAM 利用策略从而达到对核内 SRAM 的高效利用。PTO 指令抽象定义了一组抽象的 Parallel Tile Operation (PTO) InstructionPTO 指令开发者编写的程序都会转换为一组 PTO 指令。通过在不同硬件平台实现对应 PTO 指令可自动完成算子的跨平台迁移不再需要手动适配。MPMD 运行时提供了一套 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 设置无需感知复杂的内存、调度、硬件等底层实现细节对于开发者更加友好开发难度大幅下降。这一点在仓库源码中体现得非常直观。以 Indexer Prolog 算子的 Kernel 实现 quant_lightning_indexer_prolog.cpp 为例开发者通过Matrix::Matmul、Cast、Mul、RowSumSingle、Concat、Transpose、ScatterUpdate等高层 API 直接描述计算流程并在关键计算段通过TileShape::Current().SetCubeTile(...)与SetVecTile(...)指定 Tile 切分无需手工编写数据搬运指令config::SetSemanticLabel(Query-Linear); TileShape::Current().SetCubeTile({qLinear[L0M_INDEX], qLinear[L1M_INDEX]}, {qLinear[L0K_INDEX], qLinear[L1K_INDEX]}, {qLinear[L0N_INDEX], qLinear[L1N_INDEX]}, true); auto qS32 Matrix::Matmulfalse, false(DT_INT32, qNorm, wQb); // (tTile, headNum * headDim)Parallel Tile OperationPTO 指令在 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 进行矩阵乘最后对尾轴进行累加的计算示例。其中 A、B 的维度为64, 64W 的维度为64, 32。MulS、Add 和 RowSum 均为 Vector 操作设置其 Tile Shape 为32, 32Matmul 为 Cube 操作不进行切分。经过 Tiling 之后得到如下 Tile 计算图针对 A3 的 CV 分离架构框架会将 Matmul 单独切成一个子图Matmul 前的 MulS Add 切成一个子图Matmul 之后的 RowSum 切成一个子图。其中 MulS_Add 子图会被调用四次RowSum 子图会被调用两次每次处理不同的 Tile 块MPMD 运行时经过上一阶段的切图会得到许多小的子图这是 PyPTO 的最小执行单元每个子图会编译成可执行文件生成对应任务被分发到不同的核上执行。PyPTO 的运行时是 MPMD 的整个运行过程可分为如下几步根据子图间的数据依赖创建任务间依赖调度 AICPU 根据任务的拓扑顺序下发无前置依赖的任务到空闲 AI 核上AI 核收到任务后立即异步执行任务执行完成后调度 AICPU 收到回调信息解除后续任务的依赖并继续尝试下发。相比传统的 SPMD 调度要求同一时刻执行相同程序MPMD 运行时会在实际运行时根据解依赖情况调度任务执行当一个任务解依赖完成后就可以和在不同 AI 核上的任务异步并行执行从而消除不必要的全局同步开销。在 lightning_indexer_topk.cpp 中可以看到开发者通过LOOP(INDEX_LOOP_BATCH, ...)、LOOP(INDEX_LOOP_S1, ...)、LOOP(INDEX_LOOP_MATMUL, ...)等多层动态循环组织任务配合UNROLL(unrollLength)静态展开将可并行的任务交由运行时按依赖关系分发到多核执行。Human-In-The-Loop 的调优思路从算子描述到上板执行整个过程需要经过 Tile 展开、切图、运行时调度等众多阶段其中许多阶段面对的都是 NP-Hard 问题因此期望由框架自动生成高效性能的融合算子是不现实的。为此框架提供了一套可视化的集成开发环境使开发者可以查看每个阶段的输出信息控制各阶段的运行过程并将修改后的算子重新提交到 NPU 上执行得到新的结果和性能数据。通过这一机制开发者可以快速迭代在较短的周期内开发出性能优异的算子。下图展示了 PyPTO 的算子开发和调优过程环境准备与算子工程结构PyPTO 算子工程位于仓库 ops/pypto 目录是面向昇腾 Atlas A3 系列的融合算子库当前包含 DeepseekIndexerAttention、LightningIndexer 和 LightningIndexerProlog 三个算子。其运行环境要求如下项目要求产品型号Atlas A3 系列操作系统Linux ARM镜像版本cann8.3.rc1.alpha003_pt2.6.0_dsv3.2_pypto_aarch_image:v0.1驱动版本Ascend HDK 25.2.0可通过npu-smi info确认算子工程的目录结构组织清晰每个算子均包含op_host算子原型注册、InferShape/InferDataType、Tiling 代码与op_kernelKernel 源码与 UT 测试两部分例如ops/pypto/src ├── deepseek_indexer_attention # DeepseekIndexerAttention 整网融合算子 │ ├── op_host # 原型注册 / InferShape / tiling │ └── op_kernel │ ├── deepseek_indexer_attention.cpp # 算子 kernel 源码 │ ├── dynamic_mla.cpp # MLA prolog kernel │ ├── gather_after_prolog.cpp # gather kernel │ ├── lightning_indexer_prolog.cpp # lightning indexer prolog kernel │ ├── lightning_indexer_topk.cpp # LightningIndexer kernel │ └── selected_attention.cpp # selected attention kernel ├── lightning_indexer_pto # Lightning Indexer 算子 └── lightning_indexer_prolog_pto # Lightning Indexer Prolog 算子算子编译、安装与示例执行的完整流程可参考 ops/pypto/README.md核心步骤如下# 编译算子工程代码修改后需重新执行 cd /home/code/cann-recipes-infer/ops/pypto bash ./build.sh # 安装算子包 cd /home/code/cann-recipes-infer/ops/pypto/run_pkg ./customize_ops_linux.arch.run # 编译安装 whl 包 cd /home/code/cann-recipes-infer/ops/pypto/torch_ops_extension bash ./build_and_install.sh # 执行示例算子 cd /home/code/cann-recipes-infer/ops/pypto/examples python3 test_deepseek_indexer_attention.py python3 test_lightning_indexer_pto.py python3 test_lightning_indexer_prolog.py编译成功时会输出Self-extractable archive customize_ops_linux.arch.run successfully created.提示信息并在run_pkg目录下生成自定义算子包。示例脚本 test_lightning_indexer_pto.py 展示了完整的调用方式通过torch.ops.custom_pypto.npu_lightning_indexer_pto(query..., key..., weights..., actual_seq_lengths_key..., block_table..., sparse_count2048, sparse_mode3)接口完成推理调用。此外Lightning Indexer 算子已支持集成到 DeepSeek-V3.2-Exp 整网只需在网络配置 yaml 文件的model_config中增加enable_pypto: True即可开启。融合算子实例一Indexer PrologIndexer Prolog 是 DeepSeek-V3.2-Exp 的 Deepseek Indexer Attention 计算的前半部分其对应的 PyPTO 源代码位于 ops/pypto/src/lightning_indexer_prolog_pto/op_kernel计算流程图如下量化策略与计算公式Indexer Prolog 的量化策略为Q_b_proj 使用 W8A8 量化其他 Linear 均不量化Indexer Q 使用 A8 量化Indexer Cache 使用 C8 量化反量化因子以 FP16 存储同时在量化前会进行 Hadamard 变换Hadamard 变换在业界最新的低 bit 量化方案中十分常用可以有效消除 outlier。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 计算使用。源码级印证上述三段计算流在 quant_lightning_indexer_prolog.cpp 中均有清晰对应Q 路径由Matrix::MatmulInt8 输入计算 S32 中间结果→ 反量化Mul乘qNormScale与wQbScale→ RoPEQuantRope3D仅作用于头维度的 rope 部分→Matrix::BatchMatmul右乘hadamardQ→PrologQuant动态量化构成K 路径额外包含QuantLayerNormLayerNorm 防除 0 系数由attrs.eps传入量化结果通过ScatterUpdate按kCacheIndex写入 PagedAttention 的 KV CachePA_BSND排布block 大小由configs.blockSize控制W 路径则在 Matmul 之后乘以1 / (sqrt(headNum) * sqrt(headDim))缩放并转 FP16 输出。该算子的 Python 侧接口为custom_pypto.npu_lightning_indexer_prolog_pto完整参数说明见 lightning_indexer_prolog.md关键输入包括token_xbfloat16hidden 状态、wq_bint8NZ 格式query 权重、wk/weights_projbfloat16NZ 格式、ln_gamma_k/ln_beta_kkey 的 LayerNorm 缩放/偏移、cos_idx_rope/sin_idx_ropeRoPE 因子、hadamard_q/hadamard_kHadamard 变换矩阵、idx_k_cache/idx_k_scale_cache/idx_k_cache_indexkey 缓存及其反量化因子缓存和更新位置、layernorm_epsilon_k防除 0 系数layout_query默认TNDlayout_key默认PA_BSND。输出为 queryint8、query_scalefloat16与 weightsfloat16三个 Tensor。泳道图与性能分析当前该融合算子在典型 Shape4 BatchMTP1KV-Cache 64K 长度下的性能 Profiling 数据即泳道图如下图所示计算流拆解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 的调优过程将其他几个子算子进行独立性能调优达到极致性能。融合算子实例二Lightning IndexerLightning Indexer 是 DeepSeek-V3.2-Exp 的 Deepseek Indexer Attention 计算的一部分负责根据 Indexer Prolog 得到的 query、key、weights 计算索引得分其计算流程图如下Lightning Indexer 模块采用一种类似 MLP 的多查询注意力Multi-Query Attention机制来计算索引得分$$ I_{i, j} \sum_h w_h^i \cdot \text{ReLU}(\bold{q}{index, h}^i \cdot \bold{k}{index}^j) $$其中 $(w_1^i, \dots, w_{N_h}^i)^T \bold{W}_{bias}\bold{x}_i$ 表示与查询相关的、按注意力头的权重。在实际实现中$w_h^i$ 在 Indexer Prolog 模块中计算即上文中的 Indexer Weight 输出。对于每个查询 token $\bold{x}i$基于相关性得分 $I{i,j}$仅计算得分最高的前 $k$ 个缓存项所对应的索引。开发与泳道图分析流程开发者根据计算流程图进行融合算子编码开发其对应的 PyPTO 源代码位于 ops/pypto/src/lightning_indexer_pto/op_kernel/lightning_indexer_topk.cpp。基于 PyPTO 完成开发后PyPTO 提供了完善的可视化集成开发环境为开发者提供不同阶段的产物。其中最为重要的是在 NPU 上运行得到的性能 Profiling 数据即下图所示的泳道图开发者通过泳道图可以分析融合算子中各部分的耗时进而开展后续的性能调优该泳道图的横轴表示时间纵轴由一条条的泳道构成一条泳道表示一个 AI 核。每条泳道上随时间排布了对应 AI 核上执行的任务。每个任务上标记了一个标签该标签是开发者写在源码中的语义标签Semantic Label用来辅助开发者理解当前硬件在执行的任务对应到哪一段前端代码。例如在 quant_lightning_indexer_prolog.cpp 中通过config::SetSemanticLabel(Query-Linear)、Key-Hadamard、Weight-Linear等标记各计算段。选中一个任务时可以显示其前后依赖用来观察任务的执行顺序、依赖间隔等是否合理。在泳道图上可以看到不同颜色即前端开发者写的不同代码在同一时间不同计算核并行执行。上述泳道图展示的是在典型 Shape4 BatchKV-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%下一步计划经过上述优化泳道图已经比较密集。但是图中仍然有一些可优化的部分MatMul 的子图太大会导致后续 Vector 任务延迟开始第一列 Cube 任务下方仍然有大量的空白时间。通过对原始计算逻辑的分析得知这一部分可以并行因此后续可以考虑通过控制复用程度对此进行优化。最后的 TopK 是通过排序归并得出的结果可以考虑将排序部分放入前面 Vector 运算最后只做归并从而减小拖尾时间。融合算子实例三Deepseek Indexer Attention整网融合Lightning Indexer 算子展示了 PyPTO 允许开发者控制算子融合的范围。接下来通过 Deepseek Indexer Attention 的融合示例展示 PyPTO 的整网融合能力。Deepseek Indexer Attention 的计算逻辑图如下Deepseek Indexer Attention 主要由MLA Prolog、Indexer Prolog、Lightning Indexer、Sparse Flash Attention4 个相对独立的部分构成可以看出在更大的范围内融合有更大的并行度可以挖掘。PyPTO 允许开发者先手动写这独立的四个部分并通过框架将四个部分融合成一个完整的 Deepseek Indexer Attention 大算子。因此大算子的开发策略如下对于复杂融合算子可以将复杂融合算子拆分成若干个相对独立的子算子根据泳道图对各个子算子进行独立性能调优。当各子算子接近最优性能后再把子算子们进行更大程度的融合以期获得完整融合算子的最优性能。其对应的 PyPTO 源代码参考 ops/pypto/src/deepseek_indexer_attention/op_kernel/deepseek_indexer_attention.cpp其中 MLA Prolog含权重吸收技术、Gather、Lightning Indexer、Selected Attention 等子模块分别由dynamic_mla.cpp、gather_after_prolog.cpp、lightning_indexer_topk.cpp、selected_attention.cpp实现。与 Lightning Indexer 的融合算子相同也可以上板得到 Deepseek Indexer Attention 的泳道图泳道图分析大融合算子思想的优势主要在于将所有相对独立的子图融合在一起提高内存利用率并充分发挥 NPU 的性能。因此可以从图中看到DIA 整图会比单独的 Lightning 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 结果的任务停滞。改进措施及结果针对上述识别的关键性能瓶颈点逐一分析、解决开发自定义 Gather Tile OpPyPTO 允许开发者自定义 Tile Operation因此针对该场景需要开发特定的 TileOp其功能是根据 index 直接从 GM 离散搬运到 L1 的连续空间上使其能够直接进行 Matmul 的计算。这样可以省去离散转连续的额外 workspace 的申请并去除冗余的搬运性能提升非常明显。可以参考如下示意图合轴撑大 M 轴将b * s1合轴调整 Tiling撑大 M 轴用满 L0A充分提高 Matmul 的计算效率性能倍增。消除子图间 barrier消除两个子图之间的 barrier将两个子图充分融合不依赖的 task 可以提前启动可以消除流水线的气泡。消除空 reshape 子图reshape 操作不改变数据排布消除后原本依赖 reshape 的 task 的 delay 消除并减少解依赖的时间。在经过上述分析和优化后得到新版的泳道图如下可以看到使能以上优化措施后该融合算子的运行时长缩短了75%下一步计划经过上述优化泳道图中的几个关键性能瓶颈点已经被消除性能也有了比较大的提升。基于新的泳道图可以继续挖掘性能优化点持续提升性能Task 排布相对离散任务调度不够紧密需要分析任务之间的依赖消除 Pipeline Bubble。SparseAttention 中离散的数据搬运效率较低需要提高搬运的并行度。多个子图之间仍然存在较大的空隙需要再进一步的融合将有数据依赖的计算整合到同一个 LOOP 中。期望经过以上优化后泳道图中的任务会排布更加紧密。总的来说性能优化是需要分多轮迭代的每一轮解决关键的性能瓶颈点后再基于新的泳道图进行分析、持续优化。总结PyPTO 融合算子开发方法论从 DeepSeek-V3.2-Exp 三个 Indexer 系列算子的开发实践中可以提炼出如下方法论从单算子到整网融合的渐进策略对于复杂融合算子如 Deepseek Indexer Attention先拆分成若干个相对独立的子算子MLA Prolog、Indexer Prolog、Lightning Indexer、Sparse Flash Attention根据泳道图对各个子算子进行独立性能调优当各子算子接近最优性能后再进行更大程度的融合获得完整融合算子的最优性能。以泳道图为纲的调优循环泳道图是 PyPTO Human-In-The-Loop 调优的核心工具。开发者通过语义标签config::SetSemanticLabel将源码段映射到泳道图上的任务通过观察任务的执行时序、核间负载均衡、流水线气泡与拖尾快速定位瓶颈再通过调整 Tile ShapeSetCubeTile/SetVecTile、Pass 选项config::SetPassOption如L1_REUSE、自定义 Tile Operation 等手段迭代优化。显式控制融合范围与资源权衡PyPTO 允许开发者自行选定融合范围并在并行度、SRAM 利用率与数据搬运之间取得平衡如 L1 复用参数在特定 Shape 性能与动态 Shape 泛化性之间的取舍。多轮迭代、持续优化性能优化需要分多轮迭代每一轮解决关键的性能瓶颈点后再基于新的泳道图进行分析持续逼近理论极限如 Indexer Prolog 在典型 Shape 下已逼近理论极限性能Lightning Indexer 优化 27%Deepseek Indexer Attention 优化 75%。对于希望深入研究的读者建议从 ops/pypto/README.md 了解完整的环境搭建与编译流程阅读 deepseek_indexer_attention.md、lightning_indexer.md、lightning_indexer_prolog.md 三个算子文档获取完整的函数原型与参数说明并通过 ops/pypto/examples 下的示例脚本结合模型整网DeepSeek-V3.2-Exp配置enable_pypto: True进行端到端实践。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考