TPLA：高效无矢量化预填与解码推理的张量并行潜在注意力
机器学习
2025-08-26 v2 人工智能
摘要
多头潜在注意力（MLA）由 DeepSeek-V2 引入，将 key-value 状态压缩为低秩潜在向量，仅缓存该向量以减少内存。在张量并行（TP）中，注意力头跨多个设备计算，每个设备必须加载完整缓存，抵消了 MLA 相对于分组查询注意力（GQA）的优势。我们提出张量并行潜在注意力（TPLA）：一种将潜在表示和每个头的输入维度在设备之间进行分区、独立执行注意力计算，然后通过 all-reduce 合并结果的方案。TPLA 保留压缩 KV 缓存的优势，同时实现 TP 效率。与分组潜在注意力（GLA）不同，TPLA 中的每个注意力头仍可利用完整潜在表示，保持更强的表示能力。TPLA 与使用 MLA 预训练的模型兼容：支持 MLA 风格的预填，并在无需重新训练的情况下实现高效张量并行解码。通过在 TP 分片之前应用简单正交变换（如哈达德变换或 PCA），进一步减轻跨分片干扰，仅引起最小的精度退化。在保持 DeepSeek-V3 和 Kimi-K2 在常识和 LongBench 基准测试性能的同时，TPLA 将上下文长度为 32K 时的每个设备 KV 缓存降低 1.79 倍和 1.93 倍，实现了显著加速。TPLA 可与 FlashAttention-3 实现，从而实现实际的端到端加速。
引用
@article{arxiv.2508.15881,
title = {TPLA: Tensor Parallel Latent Attention for Efficient Disaggregated Prefill and Decode Inference},
author = {Xiaojuan Tang and Fanxu Meng and Pingzhi Tang and Yuxuan Wang and Di Yin and Xing Sun and Muhan Zhang},
journal= {arXiv preprint arXiv:2508.15881},
year = {2025}
}