功能列表

各功能共存情况

Y表示两功能不冲突,N表示两功能冲突。

  

控制可计算batch

多头自适应压缩

(alibi场景)

多头自适应压缩

(rope场景)

并行解码

MLA

反量化融合

qkv全量化

logN缩放

BNSD维度输入

控制可计算batch

-

Y

Y

N

Y

Y

Y

Y

Y

多头自适应压缩

(alibi场景)

-

-

N

N

N

Y

N

Y

N

多头自适应压缩

(rope场景)

-

-

-

N

N

Y

N

Y

N

并行解码

-

-

-

-

N

N

N

Y

N

MLA

-

-

-

-

-

N

N

N

N

反量化融合

-

-

-

-

-

-

N

Y

N

qkv全量化

-

-

-

-

-

-

-

N

N

logN缩放

-

-

-

-

-

-

-

-

N

BNSD维度输入

-

-

-

-

-

-

-

-

-

mask类型

mask类型不是独立特性,因paged attention算子的mask较为复杂,为便于理解,此处单独针对mask类型进行说明。

maskType

硬件类型

维度

备注

UNDEFINED

不传mask

不传mask

相当于一个全零的mask。

MASK_TYPE_NORM

Atlas 800I A2 推理产品/Atlas A2 训练系列产品

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

倒三角mask。

Atlas 推理系列产品

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

MASK_TYPE_ALIBI

Atlas 800I A2 推理产品/Atlas A2 训练系列产品

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

alibi mask。

Atlas 推理系列产品

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

MASK_TYPE_SPEC

Atlas 800I A2 推理产品/Atlas A2 训练系列产品

[num_tokens, max_seq_len]

并行解码mask。

Atlas 推理系列产品

[1, max_seq_len / 16, num_tokens, 16]

上表中Atlas 推理系列产品上max_seq_len应16对齐,且维度描述中的除法均为ceil div。

基础功能

当参数maskType、batchRunStatusEnable、quantType、calcType、compressType、scaleType均为默认值时,使用PagedAttention的基础功能,此时输入输出参数如下:

参数

维度

数据类型

格式

cpu or npu

描述

query

[num_tokens, num_head, head_size]

float16/bf16

ND

npu

各batch的query在num_tokens轴合并。

keyCache

  • Atlas 800I A2 推理产品/Atlas A2 训练系列产品:[num_blocks, block_size, kv_head_num, head_size_k];
  • Atlas 推理系列产品:[num_blocks, head_size *num_heads / 16 ,block_size, 16]

float16/bf16

  • Atlas 800I A2 推理产品/Atlas A2 训练系列产品:ND
  • Atlas 推理系列产品:NZ

npu

cache好的key。

valueCache

  • Atlas 800I A2 推理产品/Atlas A2 训练系列产品:[num_blocks, block_size, kv_head_num, head_size_v];
  • Atlas 推理系列产品:[num_blocks, head_size *num_heads / 16 ,block_size, 16]

float16/bf16

  • Atlas 800I A2 推理产品/Atlas A2 训练系列产品:ND
  • Atlas 推理系列产品:NZ

npu

cache好的value。

blockTables

[num_tokens, max_num_blocks_per_query]

int32

ND

npu

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

contextLens

[batch]

int32

ND

cpu

每个query对应的key/value的token数量。

attnOut

[num_tokens, num_head, head_size_v]

float16/bf16

ND

npu

经过计算输出的query。

控制可计算batch

多头自适应压缩(alibi场景)

多头自适应压缩(rope场景)

并行解码

MLA

反量化融合

qkv全量化

logN缩放

BNSD维度输入