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