[Requirement|需求建议]: NsaCompressAttentionInfer算子性能优化-ops-transformer-AtomGit

已关闭

[Requirement|需求建议]: NsaCompressAttentionInfer算子性能优化 #1104

dingxu创建于 3月9日关闭于 3月30日

3月9日创建

Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

Backgroud（背景信息）

本次需求对于功能无变化，意在提升importance score计算性能。
importance score计算主要有以下两个步骤

{\mathop{{p}}\nolimits_{{t}}^{{slc}}{ \left[ {j} \right] }\text{ }=\text{ }{\mathop{ \sum }\limits_{{m=0}}^{{\frac{{{l \prime }}}{{d}}-1}}{{\mathop{ \sum }\limits_{{n=0}}^{{\frac{{l}}{{d}}-1}}{\mathop{{p}}\nolimits_{{t}}^{{cmp}}{ \left[ {\frac{{{l \prime }}}{{d}}j\text{ }-\text{ }m\text{ }-\text{ }n} \right] }}}}}}\\ {\mathop{{p}}\nolimits_{{t}}^{{sl{c \prime }}}={\mathop{ \sum }\limits_{{h=1}}^{{H}}{\mathop{{p}}\nolimits_{{t}}^{{slc, \left( h \right)}}}}}

滑块计算块间和
group内再求和

当前NsaCompressAttentionInfer算子中对于importance score使用vector进行计算，会存在使用transpose，mul，add等操作，性能差。
期望通过提前生成系数矩阵，将importance score计算移入cube来提升性能

既有方案分析

分为两层循环

外循环切分行
内循环切分列
逐块计算importance score
1. 先transpose
2. 按行乘加，其中乘的系数为计算好的
3. 再transpose
4. group内reduce sum

关于系数的说明：
以l = 32, l' = 64, d = 16为例，l'/d=4，l/d=2，则 $𝑝_𝑡^{𝑠𝑙𝑐}$ [1]由以下索引的数值组成
{3A2BA8DD-0C01-4227-88D1-80201D41B2E5}.png

Benefit / Necessity （价值/作用）

通过矩阵方式计算importance score，提升性能明显。
优化前

优化后

以模型典型shape（q head num=16，kv head num=1，select block size = 64, compress block size = 32, compress stride = 16，paged block size = 128）为例，测试几组场景，单算子性能提升明显。

batch	seq length	修改前(us)	修改后(us)	耗时降低
15	1024	77	54	29.9%
15	2048	142	87	38.7%
10	4096	225	113	49.8%

Design（设计方案）

参数矩阵的生成

本版本代码实现，为了不修改接口，使用kernel内部生成参数矩阵的方式

固定参数矩阵为128行，列数32bytes对齐
额外提供validCol表示实际列数，及prefixCol表示前缀列数。以l = 32, l' = 64, d = 16为例，参数矩阵为

128行，48列（由33列pad到48列）
validCol为32
prefixCol为1

矩阵乘方案

通过矩阵乘优化importance score计算
计算流程

def ImportanceScoreCube: 
    Init result workspace to 0
    DataCopy W from GM to L1
    LoadData W from L1 to L0B                           # 参数右矩阵常驻L0B 
    for startRow in range(0, rows, rowSplit):
        endRow = startRow + rowSplit 
        for l1StartCol in range(0, cols, l1ColSplit):
            l1EndCol = l1StartCol + l1StartCol 
            DataCopy P[startRow: endRow, l1StartCol: l1EndCol] from GM to L1
            for l0startCol in range(0, l1ColSplit, l0ColSplit):
                LoadData from L1 to L0A                   # 左矩阵的分块 
                execute PW matmul
                set atomic add true
                fixpipe L0C to GM
                set atomic add false