在MLA场景下的并行解码功能。
参数 |
维度 |
数据类型 |
格式 |
cpu or npu |
描述 |
使用场景 |
|---|---|---|---|---|---|---|
query |
[num_tokens, num_head, head_size] |
float16/bf16 |
ND |
npu |
各batch的query在num_tokens轴合并。 |
基础场景 |
keyCache |
[num_blocks, block_size, kv_head_num, head_size_k] |
float16/bf16 |
ND |
npu |
kcache。 |
基础场景 |
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数量。 |
基础场景 |
mask |
[num_tokens(合轴) ，max_seqlen] |
float16/bf16 |
ND |
npu |
当maskType不为UNDEFINED时输入。 |
mask场景 |
qSeqLens |
[batch] |
int32 |
ND |
cpu |
当开启并行解码功能时需要传此tensor，每个batch对应的seqLen。 |
并行解码场景 |
参数 |
维度 |
数据类型 |
格式 |
cpu or npu |
描述 |
使用场景 |
|---|---|---|---|---|---|---|
attnOut |
[num_tokens, num_head, head_size_v] |
float16/bf16 |
ND |
npu |
经过计算输出的query。 |
基础场景 |