下载
中文
注册

AttentionFusion

功能介绍

Attention是推荐模型中会使用的典型子结构,该结构由多个算子组成。在实际网络中会出现调度时延和额外内存操作,导致整网耗时较长。Rec SDK提供了AttentionFusion融合算子,针对序列较长、key张量shape尾轴32字节对齐场景进行性能优化。

用法

使用步骤如下:

  1. 根据检查硬件条件是Rec SDK支持的产品,且CANN为8.0.RC1及之后版本,能够成功运行NPU计算网络。
  2. 安装AttentionFusion算子和AttentionFusionGrad算子,安装方法见AttentionFusion算子AttentionFusionGrad算子
  3. 通过Tensorflow框架调用AttentionFusion算子,参考样例见链接

约束说明

  • 融合算子在key的尾轴维度为8的倍数、维度较长如大于500时计算性能较好。
  • 融合算子的性能提升适用于小算子间free时间较长的情况。
  • 输入的数据:
    • 只支持3维。
    • 支持的数据类型为float32。
    • batch size均相等, 且值在(0, 2000]范围内。
    • 满足Attention的公式,shape支持对应的matmul计算。
    • 除batch size外,所有的维度满足(0, 1000]。