MLA Parallel Decoding
Description
Performs parallel decoding for MLA.
How to Enable
- Set calcType to CALC_TYPE_SPEC.
- mlaVHeadSize > 0
Input
Parameter |
Dimension |
Data Type |
Format |
cpu or npu |
Description |
Use Cases |
|---|---|---|---|---|---|---|
query |
[num_tokens, num_head, head_size] |
float16/bf16 |
ND |
npu |
query of each batch is combined along the num_tokens axis. |
Basic |
keyCache |
[num_blocks, block_size, kv_head_num, head_size_k] |
float16/bf16 |
ND |
npu |
kcache. |
Basic |
blockTables |
[num_tokens, max_num_blocks_per_query] |
int32 |
ND |
npu |
Block table of kvcache of each query. The first dimension is the token index, and the second dimension indicates the block index. |
Basic |
contextLens |
[batch] |
int32 |
ND |
cpu |
Number of key/value tokens of each query. |
Basic |
mask |
[num_tokens (combined axis), max_seqlen] |
float16/bf16 |
ND |
npu |
Required when maskType is not UNDEFINED. |
Mask |
qSeqLens |
[batch] |
int32 |
ND |
cpu |
seqLen of each batch when parallel decoding is enabled. |
Parallel decoding |
Output
Parameter |
Dimension |
Data Type |
Format |
cpu or npu |
Description |
Use Cases |
|---|---|---|---|---|---|---|
attnOut |
[num_tokens, num_head, head_size_v] |
float16/bf16 |
ND |
npu |
Query output after computation. |
Basic |
Constraints
- Same constraints as those of the MLA merging input (embed size limit).
- The values in qSeqlen must be greater than or equal to 1, and they cannot be all 1 for MTP.
- maskType can be UNDEFINED or MASK_TYPE_SPEC.
- Only MQA use cases are supported in MLA.
- The value range of head_num is (0, 256], the value range of head_size and head_size_k is (0, 576], and the value range of blocksize is (0, 128].