昇腾社区首页
中文
注册

calcType=PREFIX_ENCODER, maskType=MASK_TYPE_CAUSAL_MASK时输入输出

参数

维度

数据类型

格式

cpu/npu

描述

使用场景

query

[nTokens, qHiddenSize] 或[nTokens, headNum, headSize]

float16/bf16

ND

npu

query矩阵。

  • query的shape和key、value的不需要一致。
  • PrefixCacheLen的计算由kvseqlen减seqLen得到,由于实现方案限制,需要让prefixCacheLen 128对齐。
  • headSize需要16对齐。

基础场景

key

[numBlocks, blockSize, kvHiddenSize] 或[numBlocks, blockSize, headNum, headSize]

float16/bf16

ND

npu

key矩阵。

  • key,value的HeadSize <= 128。
  • blockSize必须设为128。
  • BlockSize × HeadSize <= 128 × 128。

基础场景

value

[numBlocks, blockSize, kvHiddenSize] 或[numBlocks, blockSize, headNum, headSize]

float16/bf16

ND

npu

value矩阵。key,value的shape必须一致。

基础场景

blockTables

[batch, maxBlockNum]

int32

ND

npu

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

blocktable 中item的value < maxBlockNum

max(kvSeqLen) / blockSize <= maxblockNum

基础场景

seqLen

[batch]

int32/uint32

ND

cpu

query对应的每个batch的序列长度。

基础场景

kvSeqLen

[batch]

int32/uint32

ND

cpu

key, value对应的每个batch的序列长度。

基础场景

output

[nTokens, qHiddenSize] 或[nTokens, headNum, headSize]

float16/bf16

ND

npu

输出。维度和query一致。

基础场景