基础功能

当参数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 训练系列产品Atlas A3 推理系列产品/Atlas A3 训练系列产品:[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 训练系列产品Atlas A3 推理系列产品/Atlas A3 训练系列产品:ND
  • Atlas 推理系列产品:NZ

npu

cache好的key。

valueCache

  • Atlas 800I A2 推理产品/Atlas A2 训练系列产品Atlas A3 推理系列产品/Atlas A3 训练系列产品:[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 训练系列产品Atlas A3 推理系列产品/Atlas A3 训练系列产品: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。