MLA(多头潜在注意力架构)凭借其低秩压缩的KV缓存设计,有效解决了长上下文推理场景下显存与带宽的双重挑战,其效能已在DeepSeek-V2等先进模型中得到充分验证。然而,主流预训练模型如LLaMA、Qwen千问等大多基于GQA架构构建,相关企业已投入大量工程优化。若直接复用MLA架构则需重新训练,成本极高。为此,依托北京大学鹏程昇腾科教创新卓越中心的算力支持,北京大学人工智能研究院助理教授张牧涵团队提出创新的TransMLA转化框架,实现了无需重训即可将主流模型无缝迁移至MLA架构。

TransMLA 框架精准针对GQA向MLA迁移的核心痛点,通过四大核心技术模块实现了精准破局:其一,实现了GQA至MLA的结构映射,破解了分组KV头与MLA单一潜在表示不兼容的难题,通过特定的线性变换,将GQA分组后的K、V向量投影融合为统一的低秩潜在表示,并配以上投影矩阵保障KV信息精准恢复,筑牢迁移基础;其二,创新性提出RoRoPE技术方案,通过巧妙的适配让位置编码顺畅融入低秩压缩流程,解决了直接对RoPE应用PCA等通用降维方法可能导致的位置信息损失或模型性能下降问题;其三,通过对RoPE中相邻频率的旋转维度进行折叠与融合,在降低参数量的同时,更有效地集中和保留关键的位置信息,从而维持模型在长序列下的语义理解能力;其四,通过均衡Key和Value矩阵在压缩前的范数分布,提升联合压缩的数值稳定性,有效减少信息损失。
在技术落地过程中,昇腾AI基础软硬件平台发挥了关键支撑作用。其高效并行计算架构满足了结构映射模块的多任务协同处理需求,保障了架构迁移效率;其优化的存储与缓存体系,为FreqFold的频率信息处理、BKV-PCA的范数均衡提供了稳定硬件基础,有效提升了KV压缩的稳定性与资源利用效率,助力TransMLA核心技术平稳落地。
经过实验验证,TransMLA的转换过程性能优势明显。在将LLaMA-2-7B模型KV缓存裁剪75%后,无需额外训练,其核心性能仅轻微损失。在32K序列长度、FP16精度下,基于昇腾平台的推理速度相较业界主流GPU平台有显著提升。依托昇腾开放的生态资源,已推动TransMLA稳定支持主流模型部署,并将集成至vLLM及SGLang等高性能推理框架生态,便于用户直接部署应用,大幅降低企业长上下文模型的落地适配成本。
TransMLA与昇腾平台的协同创新,打通了主流大模型与先进MLA架构间的鸿沟,充分发挥了昇腾全栈AI软硬件的生态优势。昇腾全链路支持实现了TransMLA“零重训、低损失”的既定目标,在保留模型参数优势的前提下,降低了企业基于昇腾平台进行AI大模型升级的门槛。这一软硬件协同典范,为长上下文推理提供了基于昇腾生态的完整解决方案,有力推动了自主计算与前沿AI技术的融合,彰显了昇腾AI的核心引领作用,为大模型产业依托自主硬件降本增效提供了可行路径。
