Merging MLA and Input KV Cache

Description

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

How to Enable

mlaVHeadSize > 0

Constraints

• The MLA and input KV cache merging does not support the Atlas inference products, ALiBi mask, multi-head adaptive compression (rope scenario), multi-head adaptive compression (ALiBi scenario), parallel decoding, logN scaling, and BNSD input layout.
• After MLA and input KV cache merging is enabled, the value range of head_size of query and keyCache is (0, 576]. When the value of head_size of keyCache or valueCache is greater than 256, the value of block_size must be less than or equal to 128, and the value of mlaVHeadSize cannot be greater than that of head_size of query and keyCache.
• MLA and KV cache merging and full quantization can be both enabled. In this case, the value range of head_size of query and key is (0, 576]. The value of mlaVHeadSize cannot be greater than that of head_size of query and key. When the value of head_size of keyCache or valueCache is greater than 256, the value of block_size must be less than or equal to 128.
• The MLA merging input supports only the MQA scenario.
• 0 < batch <= 8192 on the Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products