MLA and Input KV Cache Merging

Description

In the past, the attention operator stores k and v in the graphics memory. As the model becomes larger, the input sequence becomes longer, and the KV cache becomes larger, exerting great pressure on the graphics memory capacity and causing performance bottlenecks. Therefore, MLA compresses the KV cache into a small vector to replace the original KV cache input into the fa operator for calculation.

How to Enable

• calcType= PA_ENCODER
• mlaVHeadSize > 0

Constraints

• MLA and input KV cache merging does not support Atlas inference products, ALiBi mask, multi-head adaptive compression (rope scenario), multi-head adaptive compression (ALiBi scenario), logN scaling, or BNSD input layout.
• After MLA and input KV cache merging is enabled, the value range of headSize of the key and value is (0, 576].
• mlaVHeadSize cannot be greater than headSize of the query and key.
• After MLA and input KV cache merging is enabled, the query, key, and value can be two-dimensional [nTokens, hiddenSize], three-dimensional [nTokens, headNum, headSize], or four-dimensional [batch, seqLen, headNum, headSize]. The number of dimensions of the query and key can be different.
• MLA and input KV cache merging supports full quantization. When both functions are enabled, the headSize range of the query and key is (0, 576], and the mlaVHeadSize cannot be greater than the headSize of the query and key.