MLA for Compressed kvCache

Description

Attention operators once stored k and v in the graphics memory. As models become larger and the input sequence becomes longer, kvCache scales accordingly, exerting great pressure on graphics memory capacity and causing performance bottlenecks. Therefore, Multi-Head Latent Attention (MLA) compresses the kvCache into a small vector to replace the original kvCache input into the FA operator for calculation.

How to Enable

mlaVHeadSize > 0

Constraints

• MLA for compressed kvCache does not support the ALiBi mask, Atlas inference products, multi-head adaptive compression (RoPE), multi-head adaptive compression (ALiBi), parallel decoding, logN scaling, and BNSD input layout.
• After MLA for compressed kvCache 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.
• MLA for compressed kvCache supports only MQA.
• 0 < batch ≤ 8192 on the Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products