TransMLA: Mehrkopf-Latenzaufmerksamkeit ist alles, was Sie brauchen

TransMLA: Multi-head Latent Attention Is All You Need

February 11, 2025

Autoren: Fanxu Meng, Zengwei Yao, Muhan Zhang

cs.AI

Zusammenfassung

Moderne große Sprachmodelle (LLMs) stoßen häufig auf Kommunikationsengpässe auf aktuellen Hardwareplattformen, anstatt rein rechnerische Beschränkungen zu haben. Multi-Head Latent Attention (MLA) bewältigt diese Herausforderung, indem es Low-Rank-Matrizen in den Schlüssel-Wert (KV)-Schichten verwendet, wodurch komprimierte latente KV-Zustände zwischengespeichert werden können. Dieser Ansatz reduziert die Größe des KV-Caches im Vergleich zur traditionellen Multi-Head Attention erheblich, was zu schnelleren Inferenzzeiten führt. Darüber hinaus verwendet MLA eine Up-Projektionsmatrix, um die Ausdruckskraft zu erhöhen, wobei zusätzliche Berechnungen gegen reduzierten Kommunikationsaufwand getauscht werden. Obwohl MLA Effizienz und Effektivität in Deepseek V2/V3/R1 gezeigt hat, verlassen sich viele führende Modellanbieter immer noch auf Group Query Attention (GQA) und haben keine Pläne bekannt gegeben, MLA zu übernehmen. In diesem Artikel zeigen wir, dass GQA immer durch MLA dargestellt werden kann, während der umgekehrte Fall nicht zutrifft. Um die breitere Verwendung von MLA zu fördern, stellen wir **TransMLA** vor, eine Post-Training-Methode, die weit verbreitete GQA-basierte vortrainierte Modelle (z. B. LLaMA, Qwen, Mixtral) in MLA-basierte Modelle umwandelt. Nach der Konvertierung kann das Modell zusätzliches Training durchlaufen, um die Ausdruckskraft zu steigern, ohne die Größe des KV-Caches zu erhöhen. Darüber hinaus planen wir, MLA-spezifische Beschleunigungstechniken für die Inferenz zu entwickeln, um eine geringe Latenz in transformierten Modellen beizubehalten und somit eine effizientere Destillation von Deepseek R1 zu ermöglichen.

English

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.