RelayAttentionOperation

功能

RelayAttention算法的核心创新点是通过将与系统提示对应的矩阵-向量乘法分组为矩阵-矩阵乘法,允许对于一批输入令牌从DRAM中仅读取一次系统提示的隐藏状态(键值对),从而消除现有因果注意力计算算法中处理系统提示时存在的大量冗余内存访问,在提高效率的同时保持生成质量且无需模型重新训练。

图1 RelayAttention计算原理

算子上下文

定义

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
struct RelayAttentionParam {
    int32_t headNum = 0;
    float qkScale = 1;
    int32_t kvHeadNum = 0;
    enum MaskType : int {
        MASK_TYPE_UNDEFINED = 0,     
        MASK_TYPE_NORM,       
    };
    MaskType maskType = MASK_TYPE_UNDEFINED;
    uint8_t rsv[32] = {0};
};

参数列表

成员名称

类型

默认值

描述

headNum

int32_t

0

head数量。headNum需要大于0

qkscale

float

1.0

算子tor值。

kvHeadNum

int32_t

0

kvhead数量。取值范围为[0, 8 ]。

kvHeadNum取值分为两种情况:

  • 取值为零。
  • 取值大于0,此时kvHeadNum小于等于headNum且headNum为kvHeadNum的整数倍。

maskType

MaskType

MASK_TYPE_UNDEFINED

mask类型。

rsv[32]

uint8_t

{0}

预留参数。

输入

参数

维度

数据类型

格式

CPU or NPU

描述

query

[B, qHiddenSize]

float16/bf16

ND

NPU

query矩阵。

key

[B, [S1, N, D]]

[B, [S1, N*D]]

float16/bf16

ND

CPU

不共享key矩阵,为TensorList,可选择是否合轴。

value

[B, [S1, N, D]]

[B, [S1, N*D]]

float16/bf16

ND

CPU

不共享value矩阵,为TensorList,可选择是否合轴。

keyShare

[BS, [S2, N, D]]

[BS, [S2, N*D]]

float16/bf16

ND

CPU

共享key矩阵,为TensorList,可选择是否合轴。

valueShare

[BS, [S2, N, D]]

[BS, [S2, N*D]]

float16/bf16

ND

CPU

共享value矩阵,为TensorList,可选择是否合轴。

attentionMask

-

float16/bf16

ND

NPU

预留输入tensor。

seqLen

[B]

int32

ND

CPU

qseqlen list。

kvSeqLen

[B]

int32

ND

CPU

kvseqlen list。

kvShareMap

[B]

int32

ND

CPU

batch和共享组映射关系数组,无共享组时取值为-1

kvShareLen

[BS]

int32

ND

CPU

共享组实际长度数组。

输出

参数

维度

数据类型

格式

CPU or NPU

out

[B, qHiddenSize]

float16/bf16

ND

NPU

规格说明