Input and Output When calcType Is CALC_TYPE_PREFILL

qNope

[tokenNum, headNum * embeddimV]/[tokenNum, headNum, embeddimV]

float16/bf16

ND

NPU

Input. Query without positional encoding.

qRope

[tokenNum, headNum * 64] / [tokenNum, headNum, 64]

float16/bf16

ND

NPU

Input. Query with rotational positional encoding.

kNope

[batch, maxSeqLen, kvHeadNum * embeddimV]

float16/bf16

ND

NPU

Input. Key without positional encoding.

kRope

[batch, maxSeqLen, kvHeadNum * 64]

float16/bf16

ND

NPU

Input. Key with rotational positional encoding.

vNope

[batch, maxSeqLen, kvHeadNum * embeddimV]

float16/bf16

ND

NPU

Input. Value without positional encoding.

qSeqLen

[batch]

int32

ND

CPU

Input. Query length.

kvSeqLen

[batch]

int32

ND

CPU

Input. Key length.

mask

• Not required when UNDEFINED is used
• [512, 512] when MASK_TYPE_MASK_FREE is used
• [max(qSeqLen), max(kvSeqlen)] when MASK_TYPE_SWA_NORM is used

float16/bf16

ND

NPU

Input. Attention mask.

attenOut

[tokenNum, headNum * embeddimV]/[tokenNum, headNum, embeddimV]

float16/bf16

ND

NPU

Attention output.