TransMLA: Multi-head Latent Attention Is All You Need
Modern large language models (LLMs) often encounter communication bottlenecks
on current hardware, rather than purely computational constraints. Multi-head
Latent Attention (MLA) tackles this challenge by using low-rank matrices in the
key-value (KV) layers, thereby allowing compressed latent KV states to be
cached. This approach significantly reduces the KV cache size relative to
traditional multi-head attention, leading to faster inference. Moreover, MLA
employs an up-projection matrix to increase expressiveness, trading additional
computation for reduced communication overhead. Although MLA has demonstrated
efficiency and effectiveness in Deepseek V2/V3/R1, many major model providers
still rely on Group Query Attention (GQA) and have not announced any plans to
adopt MLA. In this paper, we show that GQA can always be represented by MLA
while maintaining the same KV cache overhead, but the converse does not hold.
To encourage broader use of MLA, we introduce **TransMLA**, a post-training
method that converts widely used GQA-based pre-trained models (e.g., LLaMA,
Qwen, Mixtral) into MLA-based models. After conversion, the model can undergo
additional training to boost expressiveness without increasing the KV cache
size. Furthermore, we plan to develop MLA-specific inference acceleration
techniques to preserve low latency in transformed models, thus enabling more
efficient distillation of Deepseek R1.
67ad5b3a007d78b391946a57
Fanxu Meng
claimed_verified
2025-02-13T08:22:03.808Z
643f55d4ec817b766686438a
/avatars/0feb460432c92ab9ada0d417a7a38f6a.svg
mengfanxu
user
fxmeng
67ad5b3a007d78b391946a58
Zengwei Yao
67ad5b3a007d78b391946a59
Muhan Zhang
TransMLA: Multi-head Latent Attention Is All You Need
Modern large language models (LLMs) often encounter communication bottlenecks
on current hardware, rather than purely computational constraints. Multi-head
Latent Attention (MLA) tackles this challenge by using low-rank matrices in the
key-value (KV) layers, thereby allowing compressed latent KV states to be
cached. This approach significantly reduces the KV cache size relative to
traditional multi-head attention, leading to faster inference. Moreover, MLA
employs an up-projection matrix to increase expressiveness, trading additional
computation for reduced communication overhead. Although MLA has demonstrated
efficiency and effectiveness in Deepseek V2/V3/R1, many major model providers
still rely on Group Query Attention (GQA) and have not announced any plans to
adopt MLA. In this paper, we show that GQA can always be represented by MLA
while maintaining the same KV cache overhead, but the converse does not hold.
To encourage broader use of MLA, we introduce **TransMLA**, a post-training
method that converts widely used GQA-based pre-trained models (e.g., LLaMA,
Qwen, Mixtral) into MLA-based models. After conversion, the model can undergo
additional training to boost expressiveness without increasing the KV cache
size. Furthermore, we plan to develop MLA-specific inference acceleration
techniques to preserve low latency in transformed models, thus enabling more
efficient distillation of Deepseek R1.
TransMLA converts GQA-based models to MLA to reduce KV cache size and improve efficiency without sacrificing expressiveness.
Multi-head Latent Attention MLA low-rank matrices KV layers latent KV states up-projection matrix Group Query Attention GQA TransMLA pre-trained models LLaMA Qwen Mixtral expressiveness KV cache overhead inference acceleration distillation