背景Triton内存管理简介昇腾硬件内存层次结构昇腾内存体系采用分层的结构设计每层都有不同的容量、带宽、访问特性。以下是昇腾NPU架构版本351x的硬件架构图https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/900/programug/Ascendcopdevg/atlas_ascendc_10_00065.html。重点关注全局内存Global Memory、L2缓存L2 Cache、统一缓冲区Unified Buffer、L1缓冲区L1 Buffer和寄存器文件包括SIMT Register File和Vector Register File两个这些存储单元。图1 NPU架构版本351x硬件架构图这些存储单元共同构成了昇腾NPU的内存层次结构通过不同层级的缓存和数据搬运机制有效平衡了内存带宽和计算性能。内存优化的策略昇腾NPU的统一缓冲区Unified Buffer, UB容量有限为了避免UB Overflow并最大化访存带宽必须将庞大的输入张量Q/K/V切分为大小合适的Block。分块策略是解决昇腾硬件片上内存有限性、优化内存访问模式、适配硬件架构差异的必要技术手段是Triton-Ascend高效实现高性能算子的核心机制。其目标是将数据分解为适合片上的内存大小。在实际开发Triton-Ascend算子的过程中我们通常会通过设置合适的BLOCK_SIZE来控制一次处理的数据块大小。Triton算子开发问题内存UB Overflow内存优化的核心技术是分块。在现实中通常需要尝试多组分块参数以得到最佳性能。如果分块过小芯片性能无法得到充分发挥。如果分配内存过大由于片上内存是有限的那么会造成内存溢出即UB Overflow问题。其中的一个重要参数就是BLOCK_SIZE。BLOCK_SIZE定义了每个Block处理的元素数量例如在向量加法算子中设置BLOCK_SIZE1024表示每个Block处理1024个元素。在一般实现中BLOCK_SIZE对于算子单次处理的数据总量具有决定性作用且通常单次处理的数据量与算子运行所需的UB大小成正比。当出现Triton算子内存溢出时通过日志我们仅能知道发生了UB overflow但是具体在何处代码导致的内存异常有哪些Tensor在使用内存还是一无所知对于问题定位造成很大的困扰。因此昇腾MindStudio提供了基于编译器产物的内存使用情况分析工具通过编译器dump内存打点结合MindStudio Insight可视化分析实现UB Overflow的代码级精准定位。内存溢出案例演示数据准备与运行我们使用Triton-Ascend官方仓库中的示例代码05-layer-norm.pyhttps://gitcode.com/Ascend/triton-ascend/blob/main/third_party/ascend/tutorials/05-layer-norm.py模拟内存溢出的场景。原始代码中BLOCK_SIZE会受到MAX_FUSED_SIZE的边界影响见下方代码保证不会内存溢出。为了构造内存溢出场景我们需要更改代码逻辑使得BLOCK_SIZE可以超过MAX_FUSED_SIZE的边界。改动如下所示# Less than 64KB per feature: enqueue fused kernel- MAX_FUSED_SIZE 65536 // x.element_size()- BLOCK_SIZE min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) BLOCK_SIZE triton.next_power_of_2(N)我们构造了可稳定复现内存溢出的测试用例下面我们用large_N_8192作为例子展示我们Triton算子内存溢出时的定位方法。环境配置NPU: 昇腾910CANN: 9.5.0Triton-Ascend: 3.2.1torch_npu: 2.7.0torch: 2.7.0用例large_N_8192参数配置在05-layer-norm.py中修改main的部分if __name__ __main__:_layer_norm(128, 8192, torch.float16)用例large_N_8192运行python 05-layer-norm.py在实际执行中我们可以看到用例large_N_8192在BLOCK_SIZE 16384预估UB使用256KB)时执行失败开始抛出内存溢出的异常。图2 large_N_8192 运行报错信息下面我们用CANN提供的毕昇编译器https://www.hiascend.com/cann/bisheng工具采集Triton算子在编译过程中出现的内存溢出数据。数据采集我们取large_N_8192用例第一个出现Overflow的用例作为例子采集数据。Triton-Ascend算子的编译调用链路如下python - triton-ascend - npuir - bishengir-compiler前面我们在执行算子时已经触发了UB Overflow此时我们找到对应的Triton缓存目录。// (triton) localhost:triton-ascend/third_party/ascend/tutorials# ls .triton_cache/${lastest_cache}_layer_norm_fwd_fused.ttadapter _layer_norm_fwd_fused.ttir通过编译执行如下命令,在当前目录下会得到一个memory_info_aiv.jsonbishengir-compiler _layer_norm_fwd_fused.ttadapter --targetAscend910B --enable-auto-multi-bufferTrue --enbale-auto-bind-sub-block-True --enable-hfusion-compile-true --enable-hivm-compile-true --enable-triton-kernel-compiletrue --mlir-pirnt-ir-before-all --mlir -print-ir-alter-all --enable-memory-dispaly参数的详细说明可以见命令行帮助信息。数据展示及分析在MindStudio Insight软件https://gitcode.com/Ascend/msinsight中导入memory_info_aiv.json文件导入后如下所示图3 MindStudio Insight 导入 json 文件总览标题栏标题栏中包含了算子名称、内存类型、编译状态信息图4 标题栏显示编译情况和具体报错信息从这里可以直接看出算子编译存在问题接下来要具体分析问题原因。内存块图该图反映了随时间变化内存的使用情况。横轴是一个虚拟的执行顺序每个时间单位相当于IRIntermediate Representation文件中的一行竖轴是片上的连续的虚拟内存地址空间。图中每个色块代表一块Tensor内存它的纵向长度表示占用的内存地址空间大小它的横向长度表示在IR文件上的生命周期。在这张图上可以反映整体运行过程中保留的连续虚拟内存地址空间的峰值和碎片情况。图5 内存块图显示连续虚拟内存地址空间经过分析。large_N_8192用例的连续虚拟内存地址空间大小中超过256 KB的地方都是存在问题的其中最后面的地方保留的连续虚拟内存地址空间最大是导致内存溢出问题最严重的地方。此外其他超过256 KB的地方也是需要处理保证总体连续虚拟内存地址空间都不高于256 KB。内存详情表在内存块图上点击具体某个内存块时该图就会展示内存块的详细信息包括大小、生命周期、申请位置、是否为临时变量。这里我们需要点击超过256 KB的内存块分析导致内存溢出的具体原因即这个内存块是在哪一行代码申请的定位到内存溢出的代码位置。图6 内存详情表展示内存块的详细信息我们选择第一个超过256 KB的内存块显示这个内存块的内存详情表。我们重点关注详情表中的Life属性这个属性显示了内存块是从IR文件的哪一行申请、在哪一行释放的。这里定位到了IR文件导致内存溢出的代码行接下来开发者可以从IR文件自行推导到Triton算子的Python文件中导致的内存溢出的代码行。使用昇腾MindStudio提供的工具链可以很好定位Triton算子内存溢出问题。未来昇腾MindStudio将持续改进不断优化使用流程帮助开发者快速定位并解决Triton算子内存溢出问题提高Triton算子开发效率。