MLA输入输出

qNope

[tokenNum, headNum, 512]

float16/bf16/int8

ND

NPU

输入，无位置编码query。

cacheMode为2：int8。

qRope

[tokenNum, headNum, 64]

float16/bf16

ND

NPU

输入，旋转位置编码query。

ctKV

[blockNum, blockSize, kvHeadNum, 512]

• cacheMode为2：

  [blockNum, kvHeadNum*512/32, blockSize, 32]

• cacheMode为3：

  [blockNum, kvHeadNum*512/16, blockSize, 16]

float16/bf16/int8

ND/NZ

NPU

输入，无位置编码ctkv。

• cacheMode为2：int8，NZ。
• cacheMode为3：float16/bf16，NZ。

kRope

[blockNum, blockSize, kvHeadNum, 64]

cacheMode为2或3：

[blockNum, kvHeadNum*64 / 16 ,blockSize, 16]

float16/bf16

ND/NZ

NPU

输入，旋转位置编码k。

cacheMode为2或3：NZ。

block_tables

[batch, maxBlockNumPerQuery]

int32

ND

NPU

输入，每个query的kvcache的block映射表。

contextLens

[batch]

int32

ND

CPU

输入，每个query对应的上下文长度，kvseqlen。

mask

• MASK_TYPE_SPEC:

  [tokenNum(合轴), maxSeqLen]

• MASK_TYPE_MASK_FREE:

  [125 + 2 * qSeqLen, 128]

float16/bf16

ND

NPU

输入，注意力掩码，maskType不为0时传入。

qseqlen

[batch]

int32

ND

CPU

输入，calcType为1时传入，每个batch对应的seqLen，取值范围为[1,4]。

qkDescale

[headNum]

float

ND

NPU

输入，cacheMode为2时传入。

pvDescale

[headNum]

float

ND

NPU

输入，cacheMode为2时传入。

attenOut

[tokenNum, headNum, 512]

float16/bf16

ND

NPU

attention输出。

Ise

[tokenNum, headNum, 1]

float16/bf16/float

ND

NPU

Ise输出。

• 当calcType为CALC_TYPE_RING和CALC_TYPE_SPEC_AND_RING时需要此tensor。
• 当CalcType为CALC_TYPE_RING时，只支持MaskType为UNDEFINED。
• 当CalcType为CALC_TYPE_SPEC_AND_RING时
  • 支持MaskType为MASK_TYPE_SPEC，cacheMode为KROPE_CTKV场景。
  • 支持MaskType为UNDEFINED，cacheMode为INT8_NZCACHE场景。
• cacheMode=INT8_NZCACHE时，dtype等于float。