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, MLA compresses the kvCache into a small vector to replace the original kvCache input into the FA operator for calculation.

How to Enable

• calcType= PA_ENCODER
• mlaVHeadSize > 0

Constraints

• MLA for compressed kvCache does not support Atlas inference products, ALiBi mask, multi-head adaptive compression (RoPE), multi-head adaptive compression (ALiBi), logN scaling, or BNSD input layout.
• After MLA for compressed kvCache 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 for compressed kvCache is enabled, the query, key, and value can be passed as a 2D tensor with [nTokens, hiddenSize], 3D tensor with [nTokens, headNum, headSize], or 4D tensor with [batch, seqLen, headNum, headSize]. The number of dimensions of the query and key can be different.
• MLA for compressed kvCache supports full quantization. When both functions are enabled, the headSize range of the query and key is (0, 576], and mlaVHeadSize cannot be greater than the headSize of the query and key.