定制融合算子

基本原理

将多个小算子替换为一个融合算子,即可以减少小算子间流通的临时内存,也可以减少需保存至反向结束的激活值内存,典型的融合算子如Flash Attention。Flash attention是Attention层内存消减的业界主流解法,它通过小块数据驻留缓存连续计算,避免了大矩阵数据在片上内存和缓存间多次交换,不需要为大矩阵开辟内存。

使用场景

对比算子级内存申请有差异,与GPU对比,算子下发及内存申请不一致。

操作步骤

  1. 采集内存profiling数据,使用性能比对工具生成对比结果。
  2. 分析比对结果中,GPU使用融合算子,NPU使用一堆小算子的差异点。
  3. 通过开发融合算子消减差异。