MLA输入输出
参数 |
维度 |
数据类型 |
格式 |
cpu/npu |
描述 |
|---|---|---|---|---|---|
qNope |
[num_tokens, num_heads, 512] |
float16/bf16/int8 |
ND |
NPU |
输入，无位置编码query。 |
qRope |
[num_tokens, num_heads, 64] |
float16/bf16 |
ND |
NPU |
输入，旋转位置编码query。 |
ctKV |
[num_blocks , block_size, kv_heads, 512] |
float16/bf16/int8 |
ND/NZ |
NPU |
输入，无位置编码ctkv。
|
kRope |
[num_blocks , block_size, kv_heads, 64] cacheMode为2或3： [blockNum, kv_heads*64 / 16 ,block_size, 16] |
float16/bf16 |
ND/NZ |
NPU |
输入，旋转位置编码k。 cacheMode为2或3：NZ。 |
block_tables |
[batch, max_num_blocks_per_query] |
int32 |
ND |
NPU |
输入，每个query的kvcache的block映射表。 |
contextLens |
[batch] |
int32 |
ND |
CPU |
输入，每个query对应的上下文长度，kvseqlen。 |
mask |
float16/bf16 |
ND |
NPU |
输入，注意力掩码，maskType不为0时传入。 |
|
qseqlen |
[batch] |
int32 |
ND |
CPU |
输入，calcType为1时传入，每个batch对应的seqLen，取值范围为[1,4]。 |
qkDescale |
[num_heads] |
float |
ND |
NPU |
输入，cacheMode为2时传入。 |
pvDescale |
[num_heads] |
float |
ND |
NPU |
输入，cacheMode为2时传入。 |
attenOut |
[num_tokens, num_heads, head_size_vo] |
float16/bf16 |
ND |
NPU |
attention输出。 |
Ise |
[num_tokens, num_heads, 1] |
float16/bf16 |
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场景。 |