将GQA模型转为多头注意力结构,推理速度提升10.6倍。
TransMLA: Multi-Head Latent Attention Is All You Need
- 用新框架直接转换GQA模型为MLA结构,兼容DeepSeek优化
- 压缩93%键值缓存,8K上下文下推理提速10.6倍,质量无损
- 仅需60亿词微调即可恢复原性能,适合迁移部署
本文提出TransMLA,一种可将任意基于GQA的预训练模型无缝转换为基于MLA的模型的框架。该方法实现与DeepSeek代码库的直接兼容,使模型能充分利用vLLM和SGlang等深度优化。通过压缩LLaMA-2-7B模型93%的键值缓存,TransMLA在8K上下文长度下实现10.6倍的推理加速,同时保持有意义的输出质量。此外,仅需60亿个词元进行微调,模型即可在多个基准测试中达到与原始模型相当的性能。当结合DeepSeek的先进特性(如FP8量化和多标记预测)时,可进一步实现更显著的推理加速。TransMLA为迁移GQA模型至MLA结构提供了实用解决方案。
原文摘要 · Abstract (English)
In this paper, we present TransMLA, a framework that seamlessly converts any GQA-based pre-trained model into an MLA-based model. Our approach enables direct compatibility with DeepSeek's codebase, allowing these models to fully leverage DeepSeek-specific optimizations such as vLLM and SGlang. By compressing 93% of the KV cache in LLaMA-2-7B, TransMLA achieves a 10.6x inference speedup at an 8K context length while preserving meaningful output quality. Additionally, the model requires only 6 billion tokens for fine-tuning to regain performance on par with the original across multiple benchmarks. TransMLA offers a practical solution for migrating GQA-based models to the MLA structure. When combined with DeepSeek's advanced features, such as FP8 quantization and Multi-Token Prediction, even greater inference acceleration can be realized.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。