TPLA：用于高效 disaggregated 预填充与解码推理的张量并行潜在注意力

在 DeepSeek-V2 中引入的多头潜在注意力 (MLA) 将键值状态压缩成一个低秩潜在向量，仅缓存该向量以减少内存。然而，在张量并行 (TP) 中，注意力头在多个设备上计算，每个设备都必须加载完整的缓存，从而侵蚀了 MLA 相对于分组查询注意力 (GQA) 的优势。我们提出了张量并行潜在注意力 (TPLA)：一种将潜在表示和每个头部的输入维度跨设备进行分区、对每个分片独立执行注意力，然后通过 all-reduce 组合结果的方案。TPLA 保留了压缩 KV 缓存的好处，同时释放了 TP 的效率。与分组潜在注意力 (GLA) 不同，TPLA 中的每个头部仍然利用完整的潜在表示，从而保持了更强的表示能力。TPLA 与使用 MLA 预训练的模型兼容：它支持 MLA 风格的预填充，并实现了高效的张量并行解码，而无需重新训练。在 TP 切片之前应用简单的正交变换（例如，Hadamard 变换或 PCA）可以进一步减轻跨分片干扰，从而实现最低的准确率下降。通过减少 DeepSeek-V3 和 Kimi-K2 的每个设备的 KV 缓存，我们在 32K 令牌上下文长度下分别实现了 1.79 倍和 1.93 倍的速度提升，同时保持了在常识和 LongBench 基准测试上的性能。TPLA 可以使用 FlashAttention-3 实现，从而能够实现实际的端到端加速。