MLA并行解码
功能说明
在MLA场景下的并行解码功能。
开启方式
- “calcType”置为CALC_TYPE_SPEC
- mlaVHeadSize > 0
输入
参数 |
维度 |
数据类型 |
格式 |
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。 |
基础场景 |
约束说明
- 与MLA合并输入的约束相同（embed大小限制）。
- qSeqlen中的值大于等于1，MTP需要不全为1。
- maskType支持UNDEFINED或者MASK_TYPE_SPEC。
- MLA下仅支持MQA场景。
- head_num范围为(0,256]，head_size、head_size_k范围为(0,576]，blocksize范围为(0,128]。
父主题： 功能列表