下载
中文
注册

PagedAttentionOperation

功能

PagedAttention。

一个Q有多个token,一个token对应多个KV的token,以token0为例,block_table代表其对应的KV的block_id,-1代表截止, 所以第二行和第四行为其目标block,context_lens则表示KV有多少个token,则代表仅有block_id为(3,4,5,9,10)是需要与Q进行计算的。

约束

  • 因硬件限制,embedding * blockSize ≤ 128 * 128。
  • block_size %16 == 0 推荐 block_size = 128,blockTables中元素的值须在[0, num_blocks)之间,head_size 范围(0, 256] 。
  • query、keyCache、valueCache、maskTensor四个入参在Atlas 推理系列产品上 只支持float16。
  • Atlas 推理系列产品上 0<batch<=2000。
  • 关于选取何种mask,如果原模型非alibi,且mask没有做padding,可以使用UNDEFINED格式的mask,即不传mask;否则依据原生模型的配置,决定是NORM还是ALIBI。

定义

struct PagedAttentionParam {
    enum MaskType : int {
        UNDEFINED = 0,
        MASK_TYPE_NORM,
        MASK_TYPE_ALIBI,
        MASK_TYPE_SPEC
    };
    enum QuantType : int {
        TYPE_QUANT_UNDEFINED = 0,
        TYPE_DEQUANT_FUSION
    };
    enum CompressType : int{
        COMPRESS_TYPE_UNDEFINED = 0,
        COMPRESS_TYPE_KVHEAD
    };
    enum CalcType : int {
        CALC_TYPE_UNDEFINED = 0,
        CALC_TYPE_SPEC
    };
    int32_t headNum = 0;
    float qkScale = 1.0; // scale after qkbmm
    int32_t kvHeadNum = 0;
    MaskType maskType = UNDEFINED; 
    bool batchRunStatusEnable = false;
    QuantType quantType = TYPE_QUANT_UNDEFINED;
    bool hasQuantOffset = false;
    CompressType compressType = COMPRESS_TYPE_UNDEFINED;
    CalcType calcType = CALC_TYPE_UNDEFINED;
};

成员

成员名称

描述

maskType

MaskType的类型值。。

  • UNDEFINED:默认值,全0的mask。
  • MASK_TYPE_NORM:倒三角mask。
  • MASK_TYPE_ALIBI:alibi mask。
  • MASK_TYPE_SPEC:并行解码mask。

quantType

量化类型。

  • TYPE_QUANT_UNDEFINED:默认值,不与量化融合。
  • TYPE_DEQUANT_FUSION:与反量化融合,只支持Atlas 800I A2推理产品

compressType

压缩类型。

  • COMPRESS_TYPE_UNDEFINED:默认值,不压缩。
  • COMPRESS_TYPE_KVHEAD:压缩key_cache, value_cache的kvHead维度, 当前仅支持Atlas 800I A2推理产品

calcType

CalcType的类型值。

  • CALC_TYPE_UNDEFINED:默认值,不开启并行解码。
  • CALC_TYPE_SPEC:并行解码功能。

headNum

query 头大小。

qkScale

算子tor值, 在Q*K^T后乘。

kvHeadNum

kv头数量。

maskType

mask类型。

batchRunStatusEnable

是否开启动态batch。

quantType

量化类型。

hasQuantOffset

开启量化功能后是否使用offset。

compressType

压缩方式。

calcType

计算类型。

输入

参数

维度

数据类型

格式

设备

描述

query

[num_tokens, num_heads, head_size]

float16/bfloat16

ND

npu

各batch的query在num_tokens轴合并。

keyCache

Atlas 800I A2推理产品:[num_blocks, block_size, k_head_num, head_size]

Atlas 推理系列产品:[num_blocks, num_tokens *num_heads / 16 ,block_size, 16]

float16/bfloat16/int8

Atlas 800I A2推理产品:ND

Atlas 推理系列产品:NZ

npu

已cache的key,当开启反量化融合功能时,只支持int8。

valueCache

Atlas 800I A2推理产品:[num_blocks, block_size, v_head_num, head_size]

Atlas 推理系列产品:[num_blocks, num_tokens *num_heads / 16 ,block_size, 16]

float16/bfloat16/int8

Atlas 800I A2推理产品:ND

Atlas 推理系列产品:NZ

npu

已cache的value,当开启反量化融合功能时,只支持int8。

blockTables

[num_tokens, max_num_blocks_per_query]

int32

ND

npu

每个query的kvcache的block table,第一维是token索引,第二维表示block索引。

contextLens

[num_tokens]

int32

ND

cpu

每个query对应的key/value的token数量(token数量 // block_size == block_num + 1)。

mask

表1

float16/bfloat16

Atlas 800I A2推理产品:ND

Atlas 推理系列产品:NZ

npu

Attention mask。

Atlas 800I A2推理产品:当maskType不为UNDEFINED时输入。

batchRunStatus

[batch]

int32

ND

cpu

当开启动态batch功能时需要传此tensor以作为哪些batch参与计算的标志位。

kDescale

[k_head_num*head_size]

int64/fp32

ND

npu

当量化类型为反量化时,步长tensor。数据类型需与vDescale保持一致。

kOffset

[k_head_num*head_size]

int32

ND

npu

当量化类型为反量化,且hasQuantOffset为true时,输入k的偏移量。

vDescale

[v_head_num*head_size]

int64/fp32

ND

npu

当量化类型为反量化时,步长tensor。

vOffset

[v_head_num*head_size]

int32

ND

npu

当量化类型为反量化,且hasQuantOffset为true时,输入v的偏移量。

表1 mask配置描述

maskType

硬件类型

维度

UNDEFINED

不传mask

不传mask

MASK_TYPE_NORM

Atlas 800I A2推理产品

[batch, 1, max_seq_len] 或 [1, max_seq_len] 或 [max_seq_len, max_seq_len]

Atlas 推理系列产品

[batch, max_seq_len / 16, 16, 16] 或 [1, max_seq_len / 16, 16, 16]

MASK_TYPE_ALIBI

Atlas 800I A2推理产品

[batch, num_head, 1, max_seq_len] 或 [num_head, 1, max_seq_len]

Atlas 推理系列产品

[batch * num_head, max_seq_len / 16, 16, 16] 或 [num_head, max_seq_len / 16, 16, 16]

MASK_TYPE_SPEC

Atlas 800I A2推理产品

[num_tokens, max_seq_len]

Atlas 推理系列产品

[1, max_seq_len / 16, num_tokens, 16]

输出

参数

维度

数据类型

格式

设备

描述

attnOut

[num_tokens, num_head, head_size]

float16/bfloat16

ND

npu

经过计算输出的query。