TransMLA:多头潜在注意力就是唯一需要的
机器学习
2025-06-13 v5 人工智能
摘要
本文提出TransMLA框架,可无缝将任何基于GQA的预训练模型转换为MLA-based模型。我们的做法使这些模型能够直接兼容DeepSeek的代码库,充分利用DeepSeek特定的优化措施,如vLLM和SGlang。通过压缩LLaMA-2-7B中93%的KV缓存,TransMLA在8K上下文长度下实现10.6倍的推理加速,同时保持有意义的输出质量。此外,该模型仅需60亿个token的微调即可恢复与原始模型在多个基准测试中相当的性能。TransMLA为将GQA-based模型迁移到MLA结构提供了实用解决方案。当与DeepSeek的高级功能结合使用时,例如FP8量化和Multi-Token Prediction,甚至可以实现更大的推理加速。
引用
@article{arxiv.2502.07864,
title = {TransMLA: Multi-Head Latent Attention Is All You Need},
author = {Fanxu Meng and Pingzhi Tang and Xiaojuan Tang and Zengwei Yao and Xing Sun and Muhan Zhang},
journal= {arXiv preprint arXiv:2502.07864},
year = {2025}
}
备注
https://github.com/fxmeng/TransMLA